refactor: setAbortController

src/background/apis/custom-api.mjs

@@ -9,7 +9,7 @@ import { getUserConfig, maxResponseTokenLength } from '../../config/index.mjs'
 import { fetchSSE } from '../../utils/fetch-sse'
 import { getConversationPairs } from '../../utils/get-conversation-pairs'
 import { isEmpty } from 'lodash-es'
-import { pushRecord } from './shared.mjs'
+import { pushRecord, setAbortController } from './shared.mjs'
 
 const getCustomApiPromptBase = async () => {
   return `I am a helpful, creative, clever, and very friendly assistant. I am familiar with various languages in the world.`
@@ -23,20 +23,7 @@ const getCustomApiPromptBase = async () => {
  * @param {string} modelName
  */
 export async function generateAnswersWithCustomApi(port, question, session, apiKey, modelName) {
-  const controller = new AbortController()
-  const stopListener = (msg) => {
-    if (msg.stop) {
-      console.debug('stop generating')
-      port.postMessage({ done: true })
-      port.onMessage.removeListener(stopListener)
-      controller.abort()
-    }
-  }
-  port.onMessage.addListener(stopListener)
-  port.onDisconnect.addListener(() => {
-    console.debug('port disconnected')
-    controller.abort()
-  })
+  const { controller, messageListener } = setAbortController(port)
 
   const prompt = getConversationPairs(session.conversationRecords, true)
   prompt.unshift({ role: 'system', content: await getCustomApiPromptBase() })
@@ -77,14 +64,11 @@ export async function generateAnswersWithCustomApi(port, question, session, apiK
     },
     async onStart() {},
     async onEnd() {
-      port.onMessage.removeListener(stopListener)
+      port.onMessage.removeListener(messageListener)
     },
     async onError(resp) {
-      port.onMessage.removeListener(stopListener)
+      port.onMessage.removeListener(messageListener)
       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}`)
     },