Bug 2011774 - Normalize tokens.existing_memories into message.memoriesApplied and persist token data r=Mardak,ai-frontend-reviewers

Responses currently include both memoriesApplied and tokens (including tokens.existing_memories). We should treat message.memoriesApplied as the single source of truth for the UI by normalizing/merging tokens.existing_memories into message.memoriesApplied during response handling (ideally before dispatching to the child so the child stays dumb). Persist the token-derived memory data so reloaded/continued conversations have the same memoriesApplied state without needing stream tokens.

Differential Revision: https://phabricator.services.mozilla.com/D280096

Author: nickgrato

browser/components/aiwindow/ui/components/ai-window/ai-window.mjs

@@ -215,6 +215,25 @@ export class AIWindow extends MozLitElement {
     this.#updateConversation();
   }
 
+  /**
+   * Processes tokens from the AI response stream and updates the message.
+   * Adds all tokens to their respective arrays in the tokens object and
+   * builds the memoriesApplied array for existing_memory tokens.
+   *
+   * @param {Array<{key: string, value: string}>} tokens - Array of parsed tokens from the stream
+   * @param {ChatMessage} currentMessage - The message object being updated
+   */
+  handleTokens = (tokens, currentMessage) => {
+    tokens.forEach(({ key, value }) => {
+      currentMessage.tokens[key].push(value);
+
+      // Build Applied Memories Array
+      if (key === "existing_memory") {
+        currentMessage.memoriesApplied.push(value);
+      }
+    });
+  };
+
   /**
    * Fetches an AI response based on the current user prompt.
    * Validates the prompt, updates conversation state, streams the response,
@@ -280,16 +299,17 @@ export class AIWindow extends MozLitElement {
           };
         }
 
+        if (!currentMessage.memoriesApplied) {
+          currentMessage.memoriesApplied = [];
+        }
+
         if (plainText) {
           currentMessage.content.body += plainText;
         }
 
         if (tokens?.length) {
-          tokens.forEach(token => {
-            currentMessage.tokens[token.key].push(token.value);
-          });
+          this.handleTokens(tokens, currentMessage);
         }
-
         this.#updateConversation();
         this.#dispatchMessageToChatContent(currentMessage);
         this.requestUpdate?.();

browser/components/aiwindow/ui/modules/ChatSql.sys.mjs

@@ -111,7 +111,9 @@ INSERT INTO message (
   jsonb(:web_search_queries_jsonb)
 )
 ON CONFLICT(message_id) DO UPDATE SET
-  is_active_branch = :is_active_branch;
+  is_active_branch = :is_active_branch,
+  memories_applied_jsonb = jsonb(:memories_applied_jsonb),
+  content_jsonb = jsonb(:content);
 `;
 
 export const CONVERSATIONS_MOST_RECENT = `

browser/components/aiwindow/ui/test/browser/browser.toml

@@ -28,6 +28,8 @@ support-files = [
 
 ["browser_aiwindow_memories.js"]
 
+["browser_aiwindow_memory_tokens.js"]
+
 ["browser_aiwindow_newtab_preload.js"]
 
 ["browser_aiwindow_search_button.js"]

browser/components/aiwindow/ui/test/browser/browser_aiwindow_memory_tokens.js

@@ -0,0 +1,92 @@
+/* Any copyright is dedicated to the Public Domain.
+ * http://creativecommons.org/publicdomain/zero/1.0/ */
+
+"use strict";
+
+add_task(async function test_handleTokens_builds_memoriesApplied_array() {
+  await SpecialPowers.pushPrefEnv({
+    set: [["browser.aiwindow.enabled", true]],
+  });
+
+  // Open AI Window to get access to ai-window element
+  const newAIWindow = await BrowserTestUtils.openNewBrowserWindow({
+    openerWindow: null,
+    aiWindow: true,
+  });
+  const browser = newAIWindow.gBrowser.selectedBrowser;
+
+  await SpecialPowers.spawn(browser, [], async () => {
+    await content.customElements.whenDefined("ai-window");
+
+    const aiWindowElement = content.document.querySelector("ai-window");
+    Assert.ok(aiWindowElement, "ai-window element should exist");
+
+    // Create a mock message that matches the real message structure
+    const mockMessage = {
+      tokens: {
+        search: [],
+        existing_memory: [],
+      },
+      memoriesApplied: [],
+    };
+
+    // Test tokens including memory and non-memory tokens
+    const testTokens = [
+      { key: "existing_memory", value: "user asked about cats" },
+      { key: "search", value: "cat behavior" }, // non-memory token
+      { key: "existing_memory", value: "user prefers tabby cats" },
+      { key: "existing_memory", value: "user has a pet cat named Fluffy" },
+    ];
+
+    // Call the actual handleTokens method
+    aiWindowElement.handleTokens(testTokens, mockMessage);
+
+    // Verify memory tokens were added to memoriesApplied
+    Assert.equal(
+      mockMessage.memoriesApplied.length,
+      3,
+      "Should have 3 memory tokens in memoriesApplied"
+    );
+    Assert.equal(
+      mockMessage.memoriesApplied[0],
+      "user asked about cats",
+      "First memory token should match"
+    );
+    Assert.equal(
+      mockMessage.memoriesApplied[1],
+      "user prefers tabby cats",
+      "Second memory token should match"
+    );
+    Assert.equal(
+      mockMessage.memoriesApplied[2],
+      "user has a pet cat named Fluffy",
+      "Third memory token should match"
+    );
+
+    // Verify all tokens (including non-memory) were added to tokens object
+    Assert.equal(
+      mockMessage.tokens.existing_memory.length,
+      3,
+      "Should have 3 existing_memory tokens in tokens object"
+    );
+    Assert.equal(
+      mockMessage.tokens.search.length,
+      1,
+      "Should have 1 search token"
+    );
+    Assert.equal(
+      mockMessage.tokens.search[0],
+      "cat behavior",
+      "Search token should match"
+    );
+
+    // Verify non-memory tokens were NOT added to memoriesApplied
+    Assert.ok(
+      !mockMessage.memoriesApplied.includes("cat behavior"),
+      "Search tokens should not be in memoriesApplied"
+    );
+  });
+
+  await BrowserTestUtils.closeWindow(newAIWindow);
+  await SpecialPowers.popPrefEnv();
+});