fix: implement URL validation and timeout handling for url like chrome webstore

Author: tonyhu-012

entrypoints/background.ts

@@ -4,6 +4,7 @@ import '@/utils/rpc'
 import { browser } from 'wxt/browser'
 import { defineBackground } from 'wxt/utils/define-background'
 
+import { INVALID_URLS } from '@/utils/constants'
 import { CONTEXT_MENU } from '@/utils/context-menu'
 import logger from '@/utils/logger'
 import { bgBroadcastRpc } from '@/utils/rpc'
@@ -36,7 +37,7 @@ export default defineBackground(() => {
 
   const setPopupStatusBasedOnUrl = async (tabId: number, url: string) => {
     const isValidUrl = /https?:\/\//.test(url ?? '')
-    if (!isValidUrl || unAttachedTabs.has(tabId)) {
+    if (!isValidUrl || unAttachedTabs.has(tabId) || INVALID_URLS.some((regex) => regex.test(url))) {
       await browser.action.setPopup({ popup: 'popup.html' })
     }
     else {

entrypoints/content/utils/chat/chat.ts

@@ -1,6 +1,7 @@
 import EventEmitter from 'events'
 import { type Ref, ref, watch } from 'vue'
 
+import { nonNullable } from '@/utils/array'
 import { parseDocument } from '@/utils/document-parser'
 import { AbortError, AppError } from '@/utils/error'
 import logger from '@/utils/logger'
@@ -296,7 +297,7 @@ export class Chat {
     const relevantTabIds = this.contextTabs.value.map((tab) => tab.tabId)
     const pages = await getDocumentContentOfTabs(relevantTabIds)
     const abortController = this.createAbortController()
-    const prompt = await nextStep(contextMsgs, pages)
+    const prompt = await nextStep(contextMsgs, pages.filter(nonNullable))
     const next = await generateObjectInBackground({
       schema: 'nextStep',
       prompt: prompt.user,
@@ -404,7 +405,7 @@ export class Chat {
     const abortController = this.createAbortController()
     const relevantTabIds = this.contextTabs.value.map((tab) => tab.tabId)
     const pages = await getDocumentContentOfTabs(relevantTabIds)
-    const prompt = await generateSearchKeywords(contextMsgs, pages)
+    const prompt = await generateSearchKeywords(contextMsgs, pages.filter(nonNullable))
     const r = await generateObjectInBackground({
       schema: 'searchKeywords',
       system: prompt.system,
@@ -459,7 +460,7 @@ export class Chat {
     }
     const relevantTabIds = this.contextTabs.value.map((tab) => tab.tabId)
     const pages = await getDocumentContentOfTabs(relevantTabIds)
-    const prompt = await chatWithPageContent(question, pages, onlineResults)
+    const prompt = await chatWithPageContent(question, pages.filter(nonNullable), onlineResults)
     await this.sendMessage(prompt.user, prompt.system, { assistantMsg: loading })
   }
 

entrypoints/content/utils/tabs.ts

@@ -1,6 +1,11 @@
+import { INVALID_URLS } from '@/utils/constants'
 import { parseDocument } from '@/utils/document-parser'
+import logger from '@/utils/logger'
 import { c2bRpc } from '@/utils/rpc'
 import { getTabStore } from '@/utils/tab-store'
+import { timeout } from '@/utils/timeout'
+
+const log = logger.child('tabs')
 
 export type TabInfo = {
   tabId: number
@@ -15,7 +20,10 @@ export type TabContent = TabInfo & {
 
 export async function getValidTabs(): Promise<TabInfo[]> {
   const tabs = await c2bRpc.getAllTabs()
-  return tabs.filter((tab) => tab.url?.startsWith('http') && tab.tabId) as TabInfo[]
+  return tabs.filter((tab) => {
+    const { url, tabId } = tab
+    return url && url.startsWith('http') && tabId && !INVALID_URLS.some((regex) => regex.test(url))
+  }) as TabInfo[]
 }
 
 export async function getTabInfo(tabId: number) {
@@ -24,7 +32,10 @@ export async function getTabInfo(tabId: number) {
 }
 
 export async function getDocumentContentOfTabs(tabIds: number[]) {
-  const contents = await Promise.all(tabIds.map((tabId) => c2bRpc.getDocumentContentOfTab(tabId)))
+  const contents = await Promise.all(tabIds.map((tabId) => timeout(c2bRpc.getDocumentContentOfTab(tabId), 5000).catch(() => {
+    log.error(`Failed to get content for tab ${tabId}, it might not be a valid HTML page or the tab is closed.`)
+    return undefined
+  })))
   return contents
 }
 

entrypoints/popup/App.vue

@@ -2,11 +2,14 @@
 import { ref } from 'vue'
 import { browser } from 'wxt/browser'
 
+import { INVALID_URLS } from '@/utils/constants'
+
 const isValidUrl = ref(false)
 browser.tabs.query({ active: true, currentWindow: true }).then(async (tabs) => {
   if (tabs.length === 1) {
     const tab = tabs[0]
-    isValidUrl.value = /https?:\/\//.test(tab.url ?? '')
+    const { url } = tab || {}
+    isValidUrl.value = !!url && /https?:\/\//.test(url) && !INVALID_URLS.some((regexp) => regexp.test(url))
   }
 })
 </script>

utils/array.ts

@@ -83,3 +83,7 @@ export class ObservableArray<T> extends Array<T> {
     this.eventEmitter.emit('change', this, oldArray)
   }
 }
+
+export function nonNullable<T>(value: T | null | undefined): value is T {
+  return value !== null && value !== undefined
+}

utils/constants.ts

@@ -5,3 +5,8 @@ export const VERSION = version.split('-')[0]
 export const OLLAMA_TUTORIAL_URL = 'https://nativemind.app/blog/tutorial/ollama-setup'
 export const OLLAMA_DOWNLOAD_URL = 'https://ollama.com/download'
 export const OLLAMA_HOMEPAGE_URL = 'https://ollama.com'
+
+export const INVALID_URLS = [
+  /^https:\/\/chromewebstore.google.com/,
+  /^https:\/\/chrome.google.com\/webstore\//,
+]

utils/error.ts

@@ -1,4 +1,4 @@
-export type ErrorCode = 'unknown' | 'requestError' | 'requestTimeout' | 'abortError'
+export type ErrorCode = 'unknown' | 'requestError' | 'requestTimeout' | 'abortError' | 'timeoutError'
 
 export abstract class AppError<Code extends ErrorCode> extends Error {
   private _appError = true
@@ -58,11 +58,23 @@ export class AbortError extends AppError<'abortError'> {
   }
 }
 
+// common timeout error for various operations
+export class TimeoutError extends AppError<'timeoutError'> {
+  constructor(message: string) {
+    super('timeoutError', message)
+  }
+
+  toLocaleMessage(): string {
+    return 'Operation timed out: ' + this.message
+  }
+}
+
 const errors = {
   unknown: UnknownError,
   requestError: ModelRequestError,
   requestTimeout: ModelRequestTimeoutError,
   abortError: AbortError,
+  timeoutError: TimeoutError,
 } satisfies Record<ErrorCode, typeof AppError<ErrorCode>>
 
 export function fromError(error: unknown): AppError<ErrorCode> {

utils/timeout.ts

@@ -0,0 +1,17 @@
+import { TimeoutError } from './error'
+
+export function timeout<R, E>(p: Promise<R>, ms: number, onTimeout?: () => E): Promise<R> {
+  return new Promise((resolve, reject) => {
+    const timer = setTimeout(() => {
+      reject(onTimeout?.() || new TimeoutError(`Promise timed out after ${ms} ms`))
+    }, ms)
+
+    p.then((result) => {
+      clearTimeout(timer)
+      resolve(result)
+    }).catch((error) => {
+      clearTimeout(timer)
+      reject(error)
+    })
+  })
+}