From f2bbc86ec9a68897cd7384469b7bffb8d99c7a18 Mon Sep 17 00:00:00 2001
From: josc146 <josStorer@outlook.com>
Date: Sat, 23 Mar 2024 22:46:50 +0800
Subject: [PATCH] improve 074b6cc9

---
 src/services/apis/chatgpt-web.mjs | 61 +++++++++++++++++++++++++++----
 src/services/apis/shared.mjs      |  8 +++-
 2 files changed, 59 insertions(+), 10 deletions(-)

diff --git a/src/services/apis/chatgpt-web.mjs b/src/services/apis/chatgpt-web.mjs
index ff8a38a..52c4a04 100644
--- a/src/services/apis/chatgpt-web.mjs
+++ b/src/services/apis/chatgpt-web.mjs
@@ -91,6 +91,47 @@ export async function getArkoseToken(config) {
   return arkoseToken
 }
 
+export async function isNeedWebsocket(accessToken) {
+  return (await request(accessToken, 'GET', '/accounts/check/v4-2023-04-27')).responseText.includes(
+    'shared_websocket',
+  )
+}
+
+export async function stopWebsocketConversation(accessToken, conversationId, wsRequestId) {
+  await request(accessToken, 'POST', '/stop_conversation', {
+    conversation_id: conversationId,
+    websocket_request_id: wsRequestId,
+  })
+}
+
+/**
+ * @type {WebSocket}
+ */
+let websocket
+/**
+ * @type {Date}
+ */
+let expired_at
+let wsCallbacks = []
+
+export async function registerWebsocket(accessToken) {
+  if (websocket && new Date() < expired_at - 300000) return
+
+  const response = JSON.parse(
+    (await request(accessToken, 'POST', '/register-websocket')).responseText,
+  )
+  if (response.wss_url) {
+    websocket = new WebSocket(response.wss_url)
+    websocket.onclose = () => {
+      websocket = null
+    }
+    websocket.onmessage = (event) => {
+      wsCallbacks.forEach((cb) => cb(event))
+    }
+    expired_at = new Date(response.expired_at)
+  }
+}
+
 /**
  * @param {Runtime.Port} port
  * @param {string} question
@@ -98,11 +139,6 @@ export async function getArkoseToken(config) {
  * @param {string} accessToken
  */
 export async function generateAnswersWithChatgptWebApi(port, question, session, accessToken) {
-  session.messageId = uuidv4()
-  if (session.parentMessageId == null) {
-    session.parentMessageId = uuidv4()
-  }
-
   let ws
   const { controller, cleanController } = setAbortController(
     port,
@@ -116,11 +152,16 @@ export async function generateAnswersWithChatgptWebApi(port, question, session,
   )
 
   const config = await getUserConfig()
-  const [models, requirementsToken, arkoseToken] = await Promise.all([
-    getModels(accessToken).catch(cleanController),
+  // eslint-disable-next-line no-unused-vars
+  const [models, requirementsToken, arkoseToken, useWebsocket] = await Promise.all([
+    getModels(accessToken).catch(cleanController), // don't throw error here
     getRequirementsToken(accessToken),
     getArkoseToken(config),
-  ])
+    isNeedWebsocket(accessToken).catch(cleanController), // don't throw error here
+  ]).catch((e) => {
+    cleanController()
+    throw e
+  })
   console.debug('models', models)
   const selectedModel = Models[session.modelName].value
   const usedModel =
@@ -149,6 +190,10 @@ export async function generateAnswersWithChatgptWebApi(port, question, session,
   let wss_url = ''
 
   const url = `${config.customChatGptWebApiUrl}${config.customChatGptWebApiPath}`
+  session.messageId = uuidv4()
+  if (session.parentMessageId == null) {
+    session.parentMessageId = uuidv4()
+  }
   const options = {
     method: 'POST',
     signal: controller.signal,
diff --git a/src/services/apis/shared.mjs b/src/services/apis/shared.mjs
index a263876..77fcab6 100644
--- a/src/services/apis/shared.mjs
+++ b/src/services/apis/shared.mjs
@@ -37,8 +37,12 @@ export function setAbortController(port, onStop, onDisconnect) {
   port.onDisconnect.addListener(disconnectListener)
 
   const cleanController = () => {
-    port.onMessage.removeListener(messageListener)
-    port.onDisconnect.removeListener(disconnectListener)
+    try {
+      port.onMessage.removeListener(messageListener)
+      port.onDisconnect.removeListener(disconnectListener)
+    } catch (e) {
+      // ignore
+    }
   }
 
   return { controller, cleanController, messageListener, disconnectListener }