improve websocket support for chatgpt web mode (#652)

Author: josStorer

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,141 +236,139 @@ 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,
     }),
   }
-  await fetchSSE(url, {
-    ...options,
-    onMessage(message) {
-      function handleMessage(data) {
-        if (data.error) {
-          throw new Error(data.error)
-        }
 
-        if (data.conversation_id) session.conversationId = data.conversation_id
-        if (data.message?.id) session.parentMessageId = data.message.id
-
-        const respAns = data.message?.content?.parts?.[0]
-        const contentType = data.message?.content?.content_type
-        if (contentType === 'text' && respAns) {
-          answer =
-            generationPrefixAnswer +
-            (generatedImageUrl && `\n\n![](${generatedImageUrl})\n\n`) +
-            respAns
-        } else if (contentType === 'code' && data.message?.status === 'in_progress') {
-          const generationText = '\n\n' + t('Generating...')
-          if (answer && !answer.endsWith(generationText)) generationPrefixAnswer = answer
-          answer = generationPrefixAnswer + generationText
-        } else if (
-          contentType === 'multimodal_text' &&
-          respAns?.content_type === 'image_asset_pointer'
-        ) {
-          const imageAsset = respAns?.asset_pointer || ''
-          if (imageAsset) {
-            fetch(
-              `${config.customChatGptWebApiUrl}/backend-api/files/${imageAsset.replace(
-                'file-service://',
-                '',
-              )}/download`,
-              {
-                credentials: 'include',
-                headers: {
-                  Authorization: `Bearer ${accessToken}`,
-                  ...(cookie && { Cookie: cookie }),
-                },
-              },
-            ).then((r) => r.json().then((json) => (generatedImageUrl = json?.download_url)))
-          }
-        }
-
-        if (answer) {
-          port.postMessage({ answer: answer, done: false, session: null })
-        }
-      }
-
-      function finishMessage() {
-        pushRecord(session, question, answer)
-        console.debug('conversation history', { content: session.conversationRecords })
-        port.postMessage({ answer: answer, done: true, session: session })
-      }
+  let answer = ''
+  let generationPrefixAnswer = ''
+  let generatedImageUrl = ''
 
-      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
+  if (useWebsocket) {
+    await registerWebsocket(accessToken)
+    const wsCallback = async (event) => {
+      let wsData
       try {
-        data = JSON.parse(message)
+        wsData = JSON.parse(event.data)
       } 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) {
+      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)
+    }
+
+    if (data.conversation_id) session.conversationId = data.conversation_id
+    if (data.message?.id) session.parentMessageId = data.message.id
+
+    const respAns = data.message?.content?.parts?.[0]
+    const contentType = data.message?.content?.content_type
+    if (contentType === 'text' && respAns) {
+      answer =
+        generationPrefixAnswer +
+        (generatedImageUrl && `\n\n![](${generatedImageUrl})\n\n`) +
+        respAns
+    } else if (contentType === 'code' && data.message?.status === 'in_progress') {
+      const generationText = '\n\n' + t('Generating...')
+      if (answer && !answer.endsWith(generationText)) generationPrefixAnswer = answer
+      answer = generationPrefixAnswer + generationText
+    } else if (
+      contentType === 'multimodal_text' &&
+      respAns?.content_type === 'image_asset_pointer'
+    ) {
+      const imageAsset = respAns?.asset_pointer || ''
+      if (imageAsset) {
+        fetch(
+          `${config.customChatGptWebApiUrl}/backend-api/files/${imageAsset.replace(
+            'file-service://',
+            '',
+          )}/download`,
+          {
+            credentials: 'include',
+            headers: {
+              Authorization: `Bearer ${accessToken}`,
+              ...(cookie && { Cookie: cookie }),
+            },
+          },
+        ).then((r) => r.json().then((json) => (generatedImageUrl = json?.download_url)))
       }
-    },
-    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}`)
-    },
-  })
+    }
+
+    if (answer) {
+      port.postMessage({ answer: answer, done: false, session: null })
+    }
+  }
+
+  function finishMessage() {
+    pushRecord(session, question, answer)
+    console.debug('conversation history', { content: session.conversationRecords })
+    port.postMessage({ answer: answer, done: true, session: session })
+  }
 }

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,