From 7543c271d0d8e1e48320fa78bd2e9f32e6f7cc77 Mon Sep 17 00:00:00 2001
From: Thomas Beverley <Thomas101@users.noreply.github.com>
Date: Wed, 13 Nov 2024 12:41:03 +0000
Subject: [PATCH] Improve window.ai typings

---
 package-lock.json                         | 8 ++++++++
 package.json                              | 1 +
 src/extension/contentscript-main/AI.ts    | 4 ++--
 src/extension/contentscript-main/index.ts | 7 +++----
 src/extension/ui-options/index.ts         | 5 +++--
 5 files changed, 17 insertions(+), 8 deletions(-)

diff --git a/package-lock.json b/package-lock.json
index 6b41587..a508d28 100644
--- a/package-lock.json
+++ b/package-lock.json
@@ -30,6 +30,7 @@
         "@babel/preset-env": "^7.26.0",
         "@stylistic/eslint-plugin": "^2.10.1",
         "@types/chrome": "^0.0.280",
+        "@types/dom-chromium-ai": "^0.0.2",
         "@types/fs-extra": "^11.0.4",
         "@types/node": "^22.9.0",
         "@types/object-path": "^0.11.4",
@@ -2771,6 +2772,13 @@
         "@types/har-format": "*"
       }
     },
+    "node_modules/@types/dom-chromium-ai": {
+      "version": "0.0.2",
+      "resolved": "https://registry.npmjs.org/@types/dom-chromium-ai/-/dom-chromium-ai-0.0.2.tgz",
+      "integrity": "sha512-5TyjYquxp5sJeNncqoQMzUvBbxZkd2s+KrzKaCwM7cSoHcG2+kpMuK8Ve04bGGWey6nL7+HQqeh9/sWDmcN6eQ==",
+      "dev": true,
+      "license": "MIT"
+    },
     "node_modules/@types/eslint": {
       "version": "9.6.1",
       "resolved": "https://registry.npmjs.org/@types/eslint/-/eslint-9.6.1.tgz",
diff --git a/package.json b/package.json
index d8a7342..9bd6036 100644
--- a/package.json
+++ b/package.json
@@ -24,6 +24,7 @@
     "@babel/preset-env": "^7.26.0",
     "@stylistic/eslint-plugin": "^2.10.1",
     "@types/chrome": "^0.0.280",
+    "@types/dom-chromium-ai": "^0.0.2",
     "@types/fs-extra": "^11.0.4",
     "@types/node": "^22.9.0",
     "@types/object-path": "^0.11.4",
diff --git a/src/extension/contentscript-main/AI.ts b/src/extension/contentscript-main/AI.ts
index 9249bfc..de8f516 100644
--- a/src/extension/contentscript-main/AI.ts
+++ b/src/extension/contentscript-main/AI.ts
@@ -18,7 +18,7 @@ class AI extends EventTarget {
   // MARK: Private
   /* **************************************************************************/
 
-  #browserAI: any
+  #browserAI: globalThis.AI | undefined
   #summarizer: AISummarizerFactory
   #writer: AIWriterFactory
   #rewriter: AIRewriterFactory
@@ -30,7 +30,7 @@ class AI extends EventTarget {
   // MARK: Lifecycle
   /* **************************************************************************/
 
-  constructor (browserAI: any) {
+  constructor (browserAI: globalThis.AI | undefined) {
     super()
 
     this.#browserAI = browserAI
diff --git a/src/extension/contentscript-main/index.ts b/src/extension/contentscript-main/index.ts
index a511132..df86e61 100644
--- a/src/extension/contentscript-main/index.ts
+++ b/src/extension/contentscript-main/index.ts
@@ -1,11 +1,10 @@
 import AI from './AI'
 
-const genericWindow = window as any
-const browserAI = genericWindow.ai
-const ai = new AI(browserAI)
+const ai = new AI(window.ai)
 
 if (process.env.BROWSER !== 'extlib') {
-  if (!genericWindow.ai) {
+  const genericWindow = window as any
+  if (!window.ai) {
     genericWindow.ai = ai
   }
   genericWindow.aibrow = ai
diff --git a/src/extension/ui-options/index.ts b/src/extension/ui-options/index.ts
index dac7249..40ff5bb 100644
--- a/src/extension/ui-options/index.ts
+++ b/src/extension/ui-options/index.ts
@@ -123,12 +123,13 @@ async function renderSettings () {
     const defaultModelList = [
       ...new Set([defaultModelId, Object.values(config.defaultModels)])
     ].map((modelId) => ({ id: modelId, name: modelId }))
-    renderDefaultModelOptions($defaultTextModelOpt, defaultModelId, defaultModelList, Boolean((window as any).ai?.assistant), useBrowserAI)
+
+    renderDefaultModelOptions($defaultTextModelOpt, defaultModelId, defaultModelList, Boolean(window.ai?.languageModel), useBrowserAI)
 
     const res = await fetch('https://aibrow.ai/api/model/list.json')
     if (res.ok) {
       const { models } = await res.json()
-      renderDefaultModelOptions($defaultTextModelOpt, defaultModelId, models, Boolean((window as any).ai?.assistant), useBrowserAI)
+      renderDefaultModelOptions($defaultTextModelOpt, defaultModelId, models, Boolean(window.ai?.languageModel), useBrowserAI)
     }
   })()