fix(firefox): implement BackgroundAliveKeeper for managing background ping intervals

tonyhu-012 · tonyhu-012 · commit 41d73deecf72 · 2025-07-11T11:20:21.000+08:00
diff --git a/entrypoints/content/utils/keepalive.ts b/entrypoints/content/utils/keepalive.ts
@@ -0,0 +1,21 @@
+import { c2bRpc } from '@/utils/rpc'
+
+// TODO: share interval id between different instances
+export class BackgroundAliveKeeper {
+  private intervalId: ReturnType<typeof setInterval> | undefined
+
+  constructor(maxTimeout = 120000) {
+    this.intervalId = setInterval(() => {
+      c2bRpc.ping()
+    }, 15000) // Check every 15 seconds
+    setTimeout(() => {
+      clearInterval(this.intervalId)
+      this.intervalId = undefined
+    }, maxTimeout)
+  }
+
+  dispose() {
+    clearInterval(this.intervalId)
+    this.intervalId = undefined
+  }
+}
diff --git a/entrypoints/content/utils/llm.ts b/entrypoints/content/utils/llm.ts
@@ -10,6 +10,8 @@ import { WebLLMSupportedModel } from '@/utils/llm/web-llm'
 import logger from '@/utils/logger'
 import { c2bRpc } from '@/utils/rpc'
 
+import { BackgroundAliveKeeper } from './keepalive'
+
 const log = logger.child('llm')
 
 const DEFAULT_PENDING_TIMEOUT = 60000 // 60 seconds
@@ -22,9 +24,11 @@ interface ExtraOptions {
 export async function* streamTextInBackground(options: Parameters<typeof c2bRpc.streamText>[0] & ExtraOptions) {
   const { abortSignal, timeout = DEFAULT_PENDING_TIMEOUT, ...restOptions } = options
   const { portName } = await c2bRpc.streamText(restOptions)
+  const aliveKeeper = new BackgroundAliveKeeper()
   const port = browser.runtime.connect({ name: portName })
   if (abortSignal) {
     abortSignal.addEventListener('abort', () => {
+      aliveKeeper.dispose()
       port.disconnect()
     })
   }
@@ -42,6 +46,7 @@ export async function* streamTextInBackground(options: Parameters<typeof c2bRpc.
         yieldData(message)
       })
       port.onDisconnect.addListener(() => {
+        aliveKeeper.dispose()
         done()
       })
       abortSignal?.addEventListener('abort', () => {
@@ -59,9 +64,11 @@ export async function* streamTextInBackground(options: Parameters<typeof c2bRpc.
 export async function* streamObjectInBackground(options: Parameters<typeof c2bRpc.streamObjectFromSchema>[0] & ExtraOptions) {
   const { abortSignal, timeout = DEFAULT_PENDING_TIMEOUT, ...restOptions } = options
   const { portName } = await c2bRpc.streamObjectFromSchema(restOptions)
+  const aliveKeeper = new BackgroundAliveKeeper()
   const port = browser.runtime.connect({ name: portName })
   if (abortSignal) {
     abortSignal.addEventListener('abort', () => {
+      aliveKeeper.dispose()
       port.disconnect()
     })
   }
@@ -79,6 +86,7 @@ export async function* streamObjectInBackground(options: Parameters<typeof c2bRp
         done()
       })
       port.onDisconnect.addListener(() => {
+        aliveKeeper.dispose()
         done()
       })
     },
@@ -92,8 +100,10 @@ export async function* streamObjectInBackground(options: Parameters<typeof c2bRp
 export async function generateObjectInBackground<S extends SchemaName>(options: Parameters<typeof c2bRpc.generateObjectFromSchema<S>>[0] & ExtraOptions) {
   const { promise: abortPromise, resolve, reject } = Promise.withResolvers<Awaited<ReturnType<typeof c2bRpc.generateObjectFromSchema<S>>>>()
   const { abortSignal, timeout = DEFAULT_PENDING_TIMEOUT, ...restOptions } = options
+  const aliveKeeper = new BackgroundAliveKeeper()
   abortSignal?.addEventListener('abort', () => {
     log.debug('generate object request aborted')
+    aliveKeeper.dispose()
     reject(new AbortError('Aborted'))
   })
   const timeoutTimer = setTimeout(() => {
@@ -111,6 +121,8 @@ export async function generateObjectInBackground<S extends SchemaName>(options:
       return result
     }).catch((error) => {
       throw fromError(error)
+    }).finally(() => {
+      aliveKeeper.dispose()
     })
   return await Promise.race([abortPromise, promise])
 }
@@ -121,6 +133,7 @@ export async function deleteOllamaModel(modelId: string) {
 
 export async function* pullOllamaModel(modelId: string, abortSignal?: AbortSignal) {
   const { portName } = await c2bRpc.pullOllamaModel(modelId)
+  const aliveKeeper = new BackgroundAliveKeeper()
   const port = browser.runtime.connect({ name: portName })
   abortSignal?.addEventListener('abort', () => {
     port.disconnect()
@@ -137,6 +150,7 @@ export async function* pullOllamaModel(modelId: string, abortSignal?: AbortSigna
       }
     })
     port.onDisconnect.addListener(() => {
+      aliveKeeper.dispose()
       done()
     })
   })
diff --git a/entrypoints/content/utils/search.ts b/entrypoints/content/utils/search.ts
@@ -4,6 +4,8 @@ import { toAsyncIter } from '@/utils/async'
 import { c2bRpc } from '@/utils/rpc'
 import { SearchingMessage } from '@/utils/search'
 
+import { BackgroundAliveKeeper } from './keepalive'
+
 export interface SearchOptions {
   resultLimit?: number
   abortSignal?: AbortSignal
@@ -15,7 +17,9 @@ export class SearchScraper {
     const { abortSignal, ...restOptions } = options || { engine: 'google' }
     const { portName } = await c2bRpc.searchOnline(queryList, restOptions)
     const port = browser.runtime.connect({ name: portName })
+    const aliveKeeper = new BackgroundAliveKeeper()
     options?.abortSignal?.addEventListener('abort', () => {
+      aliveKeeper.dispose()
       port.disconnect()
     })
     const iter = toAsyncIter<SearchingMessage>((yieldData, done) => {
@@ -29,6 +33,7 @@ export class SearchScraper {
       })
       const onDisconnect = () => {
         done()
+        aliveKeeper.dispose()
         port.onDisconnect.removeListener(onDisconnect)
       }
       port.onMessage.addListener(onMessage)
diff --git a/utils/rpc/background-fns.ts b/utils/rpc/background-fns.ts
@@ -490,10 +490,15 @@ function captureVisibleTab(windowId?: number, options?: Browser.tabs.CaptureVisi
   return screenCaptureBase64Url
 }
 
+function ping() {
+  return 'pong'
+}
+
 export const backgroundFunctions = {
   emit: <E extends keyof Events>(ev: E, ...args: Parameters<Events[E]>) => {
     eventEmitter.emit(ev, ...args)
   },
+  ping,
   getTabInfo: (_tabInfo?: { tabId: number }) => _tabInfo as { tabId: number, title: string, faviconUrl?: string, url: string }, // a trick to get tabId
   generateText,
   generateTextAsync,
