improve websocket support for chatgpt web mode (#652)

src/services/apis/chatgpt-web.mjs

@@ -97,6 +97,13 @@ export async function isNeedWebsocket(accessToken) {
   )
 }
 
+export async function sendWebsocketConversation(accessToken, options) {
+  const apiUrl = (await getUserConfig()).customChatGptWebApiUrl
+  const response = await fetch(`${apiUrl}/backend-api/conversation`, options).then((r) => r.json())
+  console.debug(`request: ws /conversation`, response)
+  return { conversationId: response.conversation_id, wsRequestId: response.websocket_request_id }
+}
+
 export async function stopWebsocketConversation(accessToken, conversationId, wsRequestId) {
   await request(accessToken, 'POST', '/stop_conversation', {
     conversation_id: conversationId,
@@ -111,11 +118,11 @@ let websocket
 /**
  * @type {Date}
  */
-let expired_at
+let expires_at
 let wsCallbacks = []
 
 export async function registerWebsocket(accessToken) {
-  if (websocket && new Date() < expired_at - 300000) return
+  if (websocket && new Date() < expires_at - 300000) return
 
   const response = JSON.parse(
     (await request(accessToken, 'POST', '/register-websocket')).responseText,
@@ -124,11 +131,13 @@ export async function registerWebsocket(accessToken) {
     websocket = new WebSocket(response.wss_url)
     websocket.onclose = () => {
       websocket = null
+      expires_at = null
+      console.debug('global websocket closed')
     }
     websocket.onmessage = (event) => {
       wsCallbacks.forEach((cb) => cb(event))
     }
-    expired_at = new Date(response.expired_at)
+    expires_at = new Date(response.expires_at)
   }
 }
 
@@ -139,15 +148,14 @@ export async function registerWebsocket(accessToken) {
  * @param {string} accessToken
  */
 export async function generateAnswersWithChatgptWebApi(port, question, session, accessToken) {
-  let ws
   const { controller, cleanController } = setAbortController(
     port,
     () => {
-      if (ws) ws.close()
+      if (session.wsRequestId)
+        stopWebsocketConversation(accessToken, session.conversationId, session.wsRequestId)
     },
     () => {
       if (session.autoClean) deleteConversation(accessToken, session.conversationId)
-      if (ws) ws.close()
     },
   )
 
@@ -184,13 +192,9 @@ export async function generateAnswersWithChatgptWebApi(port, question, session,
     ).value
   }
 
-  let answer = ''
-  let generationPrefixAnswer = ''
-  let generatedImageUrl = ''
-  let wss_url = ''
-
   const url = `${config.customChatGptWebApiUrl}${config.customChatGptWebApiPath}`
   session.messageId = uuidv4()
+  session.wsRequestId = uuidv4()
   if (session.parentMessageId == null) {
     session.parentMessageId = uuidv4()
   }
@@ -232,11 +236,90 @@ export async function generateAnswersWithChatgptWebApi(port, question, session,
       parent_message_id: session.parentMessageId,
       timezone_offset_min: new Date().getTimezoneOffset(),
       history_and_training_disabled: config.disableWebModeHistory,
+      websocket_request_id: session.wsRequestId,
     }),
   }
+
+  let answer = ''
+  let generationPrefixAnswer = ''
+  let generatedImageUrl = ''
+
+  if (useWebsocket) {
+    await registerWebsocket(accessToken)
+    const wsCallback = async (event) => {
+      let wsData
+      try {
+        wsData = JSON.parse(event.data)
+      } catch (error) {
+        console.debug('json error', error)
+        return
+      }
+      if (wsData.type === 'http.response.body') {
+        let body
+        try {
+          body = atob(wsData.body).replace(/^data:/, '')
+          const data = JSON.parse(body)
+          console.debug('ws message', data)
+          if (wsData.conversation_id === session.conversationId) {
+            handleMessage(data)
+          }
+        } catch (error) {
+          if (body && body.trim() === '[DONE]') {
+            console.debug('ws message', '[DONE]')
+            if (wsData.conversation_id === session.conversationId) {
+              finishMessage()
+              wsCallbacks = wsCallbacks.filter((cb) => cb !== wsCallback)
+            }
+          } else {
+            console.debug('json error', error)
+          }
+        }
+      }
+    }
+    wsCallbacks.push(wsCallback)
+    const { conversationId, wsRequestId } = await sendWebsocketConversation(accessToken, options)
+    session.conversationId = conversationId
+    session.wsRequestId = wsRequestId
+    port.postMessage({ session: session })
+  } else {
     await fetchSSE(url, {
       ...options,
       onMessage(message) {
+        console.debug('sse message', message)
+        if (message.trim() === '[DONE]') {
+          finishMessage()
+          return
+        }
+        let data
+        try {
+          data = JSON.parse(message)
+        } catch (error) {
+          console.debug('json error', error)
+          return
+        }
+        handleMessage(data)
+      },
+      async onStart() {
+        // sendModerations(accessToken, question, session.conversationId, session.messageId)
+      },
+      async onEnd() {
+        port.postMessage({ done: true })
+        cleanController()
+      },
+      async onError(resp) {
+        cleanController()
+        if (resp instanceof Error) throw resp
+        if (resp.status === 403) {
+          throw new Error('CLOUDFLARE')
+        }
+        const error = await resp.json().catch(() => ({}))
+        throw new Error(
+          !isEmpty(error) ? JSON.stringify(error) : `${resp.status} ${resp.statusText}`,
+        )
+      },
+    })
+  }
+
   function handleMessage(data) {
     if (data.error) {
       throw new Error(data.error)
@@ -288,85 +371,4 @@ export async function generateAnswersWithChatgptWebApi(port, question, session,
     console.debug('conversation history', { content: session.conversationRecords })
     port.postMessage({ answer: answer, done: true, session: session })
   }
-
-      console.debug('sse message', message)
-      if (message.trim() === '[DONE]') {
-        if (!wss_url) {
-          finishMessage()
-        } else {
-          ws = new WebSocket(wss_url)
-          ws.onmessage = (event) => {
-            let wsData
-            try {
-              wsData = JSON.parse(event.data)
-            } catch (error) {
-              console.debug('json error', error)
-              return
-            }
-            if (wsData.type === 'http.response.body') {
-              let body
-              try {
-                body = atob(wsData.body).replace(/^data:/, '')
-                const data = JSON.parse(body)
-                console.debug('ws message', data)
-                if (wsData.conversation_id === session.conversationId) {
-                  handleMessage(data)
-                }
-              } catch (error) {
-                if (body && body.trim() === '[DONE]') {
-                  console.debug('ws message', '[DONE]')
-                  if (wsData.conversation_id === session.conversationId) {
-                    finishMessage()
-                    ws.close()
-                  }
-                } else {
-                  console.debug('json error', error)
-                }
-              }
-            }
-          }
-          ws.onopen = () => {
-            // fetch(url, options)
-          }
-          ws.onclose = () => {
-            port.postMessage({ done: true })
-            cleanController()
-          }
-          ws.onerror = (event) => {
-            console.debug('ws error', event)
-            port.postMessage({ error: event })
-            cleanController()
-          }
-        }
-        return
-      }
-      let data
-      try {
-        data = JSON.parse(message)
-      } catch (error) {
-        console.debug('json error', error)
-        return
-      }
-      if (data.wss_url) wss_url = data.wss_url
-      handleMessage(data)
-    },
-    async onStart() {
-      // sendModerations(accessToken, question, session.conversationId, session.messageId)
-    },
-    async onEnd() {
-      if (!wss_url) {
-        port.postMessage({ done: true })
-        cleanController()
-      }
-    },
-    async onError(resp) {
-      cleanController()
-      if (resp instanceof Error) throw resp
-      if (resp.status === 403) {
-        throw new Error('CLOUDFLARE')
-      }
-      const error = await resp.json().catch(() => ({}))
-      throw new Error(!isEmpty(error) ? JSON.stringify(error) : `${resp.status} ${resp.statusText}`)
-    },
-  })
 }

src/services/init-session.mjs

@@ -16,6 +16,7 @@ import { v4 as uuidv4 } from 'uuid'
  * @property {string|null} conversationId - chatGPT web mode
  * @property {string|null} messageId - chatGPT web mode
  * @property {string|null} parentMessageId - chatGPT web mode
+ * @property {string|null} wsRequestId - chatGPT web mode
  * @property {string|null} bingWeb_encryptedConversationSignature
  * @property {string|null} bingWeb_conversationId
  * @property {string|null} bingWeb_clientId
@@ -63,6 +64,7 @@ export function initSession({
     conversationId: null,
     messageId: null,
     parentMessageId: null,
+    wsRequestId: null,
 
     // bing
     bingWeb_encryptedConversationSignature: null,