From 1ad1b7004d6fd9771aa0ab405f7aa6e9e1a95a9d Mon Sep 17 00:00:00 2001
From: Jordan Wages <jordan@jordanwages.com>
Date: Fri, 27 Jun 2025 01:05:27 -0500
Subject: [PATCH] Add cache clearing feature and update UI

---
 background.js                 | 73 ++++++++++++++++++++++++++++++++++-
 manifest.json                 |  4 +-
 modules/AiClassifier.js       | 21 +++++++++-
 resources/clearCacheButton.js | 22 +++++++++++
 4 files changed, 116 insertions(+), 4 deletions(-)
 create mode 100644 resources/clearCacheButton.js

diff --git a/background.js b/background.js
index e47bf99..de6b595 100644
--- a/background.js
+++ b/background.js
@@ -147,6 +147,37 @@ async function applyAiRules(idsInput) {
     return queue;
 }
 
+async function clearCacheForMessages(idsInput) {
+    const ids = Array.isArray(idsInput) ? idsInput : [idsInput];
+    if (!ids.length) return;
+
+    if (!aiRules.length) {
+        const { aiRules: stored } = await storage.local.get("aiRules");
+        aiRules = Array.isArray(stored) ? stored.map(r => {
+            if (r.actions) return r;
+            const actions = [];
+            if (r.tag) actions.push({ type: 'tag', tagKey: r.tag });
+            if (r.moveTo) actions.push({ type: 'move', folder: r.moveTo });
+            const rule = { criterion: r.criterion, actions };
+            if (r.stopProcessing) rule.stopProcessing = true;
+            return rule;
+        }) : [];
+    }
+
+    const keys = [];
+    for (const msg of ids) {
+        const id = msg?.id ?? msg;
+        for (const rule of aiRules) {
+            const key = await sha256Hex(`${id}|${rule.criterion}`);
+            keys.push(key);
+        }
+    }
+    if (keys.length) {
+        await AiClassifier.removeCacheEntries(keys);
+        showTransientIcon("resources/img/done.png");
+    }
+}
+
 (async () => {
     logger = await import(browser.runtime.getURL("logger.js"));
     try {
@@ -191,6 +222,21 @@ async function applyAiRules(idsInput) {
     }
 
     logger.aiLog("background.js loaded – ready to classify", {debug: true});
+    if (browser.messageDisplayAction) {
+        browser.messageDisplayAction.setTitle({ title: "Classify" });
+        if (browser.messageDisplayAction.setLabel) {
+            browser.messageDisplayAction.setLabel({ label: "Classify" });
+        }
+    }
+    if (browser.messageDisplayScripts?.registerScripts) {
+        try {
+            await browser.messageDisplayScripts.registerScripts([
+                { js: [browser.runtime.getURL("resources/clearCacheButton.js")] }
+            ]);
+        } catch (e) {
+            logger.aiLog("failed to register message display script", { level: 'warn' }, e);
+        }
+    }
 
     browser.menus.create({
         id: "apply-ai-rules-list",
@@ -202,6 +248,16 @@ async function applyAiRules(idsInput) {
         title: "Apply AI Rules",
         contexts: ["message_display_action"],
     });
+    browser.menus.create({
+        id: "clear-ai-cache-list",
+        title: "Clear AI Cache",
+        contexts: ["message_list"],
+    });
+    browser.menus.create({
+        id: "clear-ai-cache-display",
+        title: "Clear AI Cache",
+        contexts: ["message_display_action"],
+    });
 
     if (browser.messageDisplayAction) {
         browser.messageDisplayAction.onClicked.addListener(async (tab) => {
@@ -220,6 +276,10 @@ async function applyAiRules(idsInput) {
             const ids = info.selectedMessages?.messages?.map(m => m.id) ||
                          (info.messageId ? [info.messageId] : []);
             await applyAiRules(ids);
+        } else if (info.menuItemId === "clear-ai-cache-list" || info.menuItemId === "clear-ai-cache-display") {
+            const ids = info.selectedMessages?.messages?.map(m => m.id) ||
+                         (info.messageId ? [info.messageId] : []);
+            await clearCacheForMessages(ids);
         }
     });
 
@@ -243,8 +303,17 @@ async function applyAiRules(idsInput) {
             // rethrow so the caller sees the failure
             throw err;
         }
-    }
-    else {
+    } else if (msg?.type === "sortana:clearCacheForDisplayed") {
+        try {
+            const tabs = await browser.tabs.query({ active: true, currentWindow: true });
+            const tabId = tabs[0]?.id;
+            const msgs = tabId ? await browser.messageDisplay.getDisplayedMessages(tabId) : [];
+            const ids = msgs.map(m => m.id);
+            await clearCacheForMessages(ids);
+        } catch (e) {
+            logger.aiLog("failed to clear cache from message script", { level: 'error' }, e);
+        }
+    } else {
         logger.aiLog("Unknown message type, ignoring", {level: 'warn'}, msg?.type);
     }
 });
diff --git a/manifest.json b/manifest.json
index 1610453..0e9b20a 100644
--- a/manifest.json
+++ b/manifest.json
@@ -22,7 +22,9 @@
     "default_icon": "resources/img/logo32.png"
   },
   "message_display_action": {
-    "default_icon": "resources/img/logo32.png"
+    "default_icon": "resources/img/logo32.png",
+    "default_title": "Classify",
+    "default_label": "Classify"
   },
   "background": { "scripts": [ "background.js" ] },
   "options_ui": {
diff --git a/modules/AiClassifier.js b/modules/AiClassifier.js
index d2cae3a..cec4a31 100644
--- a/modules/AiClassifier.js
+++ b/modules/AiClassifier.js
@@ -210,6 +210,25 @@ function cacheResult(cacheKey, matched) {
   }
 }
 
+async function removeCacheEntries(keys = []) {
+  if (!Array.isArray(keys)) {
+    keys = [keys];
+  }
+  if (!gCacheLoaded) {
+    await loadCache();
+  }
+  let removed = false;
+  for (let key of keys) {
+    if (gCache.delete(key)) {
+      removed = true;
+      aiLog(`[AiClassifier] Removed cache entry '${key}'`, {debug: true});
+    }
+  }
+  if (removed) {
+    await saveCache();
+  }
+}
+
 function classifyTextSync(text, criterion, cacheKey = null) {
   if (!Services?.tm?.spinEventLoopUntil) {
     throw new Error("classifyTextSync requires Services");
@@ -288,4 +307,4 @@ async function classifyText(text, criterion, cacheKey = null) {
   }
 }
 
-export { classifyText, classifyTextSync, setConfig };
+export { classifyText, classifyTextSync, setConfig, removeCacheEntries };
diff --git a/resources/clearCacheButton.js b/resources/clearCacheButton.js
new file mode 100644
index 0000000..a4d2adb
--- /dev/null
+++ b/resources/clearCacheButton.js
@@ -0,0 +1,22 @@
+(function() {
+  function addButton() {
+    const toolbar = document.querySelector("#header-view-toolbar") ||
+                    document.querySelector("#mail-toolbox toolbar");
+    if (!toolbar || document.getElementById('sortana-clear-cache-button')) return;
+    const button = document.createXULElement ?
+          document.createXULElement('toolbarbutton') :
+          document.createElement('button');
+    button.id = 'sortana-clear-cache-button';
+    button.setAttribute('label', 'Clear Cache');
+    button.className = 'toolbarbutton-1';
+    button.addEventListener('command', () => {
+      browser.runtime.sendMessage({ type: 'sortana:clearCacheForDisplayed' });
+    });
+    toolbar.appendChild(button);
+  }
+  if (document.readyState === 'complete' || document.readyState === 'interactive') {
+    addButton();
+  } else {
+    document.addEventListener('DOMContentLoaded', addButton, { once: true });
+  }
+})();