From 8e0f32b444667dc1ccd98d88a5d775dc4bb02c52 Mon Sep 17 00:00:00 2001
From: Travis Fischer <fisch0920@gmail.com>
Date: Thu, 15 Dec 2022 01:35:09 -0600
Subject: [PATCH] feat: add timeout support to browser sendMessage

---
 src/chatgpt-api-browser.ts | 45 ++++++++++++++++++++++++++++----------
 1 file changed, 33 insertions(+), 12 deletions(-)

diff --git a/src/chatgpt-api-browser.ts b/src/chatgpt-api-browser.ts
index d626358..91d0e33 100644
--- a/src/chatgpt-api-browser.ts
+++ b/src/chatgpt-api-browser.ts
@@ -1,5 +1,6 @@
 import delay from 'delay'
 import html2md from 'html-to-md'
+import pTimeout from 'p-timeout'
 import type { Browser, HTTPRequest, HTTPResponse, Page } from 'puppeteer'
 
 import { getBrowser, getOpenAIAuth } from './openai-auth'
@@ -275,7 +276,14 @@ export class ChatGPTAPIBrowser {
     }
   }
 
-  async sendMessage(message: string): Promise<string> {
+  async sendMessage(
+    message: string,
+    opts: {
+      timeoutMs?: number
+    } = {}
+  ): Promise<string> {
+    const { timeoutMs } = opts
+
     const inputBox = await this._getInputBox()
     if (!inputBox) throw new Error('not signed in')
 
@@ -294,19 +302,32 @@ export class ChatGPTAPIBrowser {
       }
     }
 
-    do {
-      await delay(1000)
+    const responseP = new Promise<string>(async (resolve, reject) => {
+      try {
+        do {
+          await delay(1000)
 
-      // TODO: this logic needs some work because we can have repeat messages...
-      const newLastMessage = await this.getLastMessage()
-      if (
-        newLastMessage &&
-        lastMessage?.toLowerCase() !== newLastMessage?.toLowerCase()
-      ) {
-        await delay(5000)
-        return newLastMessage
+          // TODO: this logic needs some work because we can have repeat messages...
+          const newLastMessage = await this.getLastMessage()
+          if (
+            newLastMessage &&
+            lastMessage?.toLowerCase() !== newLastMessage?.toLowerCase()
+          ) {
+            return resolve(newLastMessage)
+          }
+        } while (true)
+      } catch (err) {
+        return reject(err)
       }
-    } while (true)
+    })
+
+    if (timeoutMs) {
+      return pTimeout(responseP, {
+        milliseconds: timeoutMs
+      })
+    } else {
+      return responseP
+    }
   }
 
   async resetThread() {