Bug 2064128 - Show the Memories Applied button for PUWYLO response r=pdahiya,ai-security-reviewers

- When memories are enabled, marks the known PUWYLO memory as applied once the response completes

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

browser/components/aiwindow/models/ConversationSuggestions.sys.mjs

@@ -911,6 +911,7 @@ export async function constructConversationToResumeActivity(
     ...(conversationId ? { id: conversationId } : {}),
     title: resumeActivitySuggestion.content.headline,
   });
+  conversation.promptEmbeddedMemories = [resumeActivitySuggestion.memory];
 
   const [
     { prompt: chatSystemPrompt },

browser/components/aiwindow/ui/components/ai-action-confirmation/ai-action-confirmation.css

@@ -4,6 +4,8 @@
 
 :host {
   display: block;
+  /* Keep descendant z-indexes within this component. */
+  isolation: isolate;
 
   @media not (forced-colors) {
     --ai-action-confirmation-accent-color: var(--smartwindow-text-color-accent-primary, light-dark(var(--color-violet-90), var(--color-violet-20)));

browser/components/aiwindow/ui/components/ai-chat-content/chat-assistant-footer/applied-memories-button.css

@@ -177,6 +177,7 @@ moz-button.memories-trigger {
   opacity: 0;
   pointer-events: none;
   position: absolute;
+  z-index: 20;
   /* stylelint-disable-next-line stylelint-plugin-mozilla/use-design-tokens */
   max-width: min(100%, var(--size-layout-large));
   width: 100%;

browser/components/aiwindow/ui/components/ai-website-confirmation/ai-website-confirmation.css

@@ -10,6 +10,8 @@
 
   display: block;
   position: relative;
+  /* Keep descendant z-indexes within this component. */
+  isolation: isolate;
   /* stylelint-disable-next-line stylelint-plugin-mozilla/use-design-tokens */
   max-width: min(100%, var(--size-layout-large));
 }

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

@@ -133,6 +133,13 @@ export class ChatConversation extends Conversation {
    */
   #pendingBrowserActionTelemetry = new Map();
 
+  /**
+   * Uncited memories embedded in the current prompt. Cleared after completion.
+   *
+   * @type {Array<object>}
+   */
+  promptEmbeddedMemories = [];
+
   /**
    * A mapping of a URL to its unique URL token. URL tokens are used as shortened
    * versions of URLs to help the model deal with very long URLs. Very long URLs are
@@ -481,12 +488,23 @@ export class ChatConversation extends Conversation {
 
     // Only resolve used memories once the entire assistant turn is complete,
     // including all tool calls
-    const citedMemoryIds = currentMessage.tokens?.existing_memory ?? [];
-    if (!result.pendingToolCalls?.length && citedMemoryIds.length) {
-      currentMessage.memoriesApplied =
-        await lazy.MemoriesManager.resolveUsedMemories(citedMemoryIds);
+    if (!result.pendingToolCalls?.length) {
+      const citedMemoryIds = currentMessage.tokens?.existing_memory ?? [];
+      const promptEmbeddedMemoryIds = this.promptEmbeddedMemories.map(
+        memory => memory.id
+      );
+      this.promptEmbeddedMemories = [];
 
-      this.emit("chat-conversation:message-update", currentMessage);
+      const memoryIds = [...citedMemoryIds, ...promptEmbeddedMemoryIds];
+      const memoriesApplied = memoryIds.length
+        ? await lazy.MemoriesManager.resolveUsedMemories(memoryIds)
+        : [];
+
+      if (memoriesApplied.length) {
+        currentMessage.memoriesApplied = memoriesApplied;
+
+        this.emit("chat-conversation:message-update", currentMessage);
+      }
     }
 
     // Drop the pending chunk write rather than flushing it: the full write

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

@@ -872,6 +872,72 @@ add_task(async function test_resume_prompt_click_injects_context() {
   }
 });
 
+add_task(async function test_resume_prompt_click_marks_memory_applied() {
+  const sb = sinon.createSandbox();
+  const { server, port } = startMockOpenAI({
+    streamChunks: ["Sure, here's an update."],
+  });
+  try {
+    await SpecialPowers.pushPrefEnv({
+      set: [["browser.smartwindow.endpoint", `http://localhost:${port}/v1`]],
+    });
+
+    try {
+      await testResumeActivityClick(
+        sb,
+        async ({ browser, buttons }) => {
+          buttons[0].click();
+
+          const aiWindowEl = browser.contentDocument.querySelector("ai-window");
+          const aichatBrowser = await TestUtils.waitForCondition(
+            () => aiWindowEl.shadowRoot?.querySelector("#aichat-browser"),
+            "Wait for aichat-browser"
+          );
+
+          await SpecialPowers.spawn(aichatBrowser, [], async () => {
+            const chatContent =
+              content.document.querySelector("ai-chat-content");
+
+            await ContentTaskUtils.waitForMutationCondition(
+              chatContent.shadowRoot,
+              { childList: true, subtree: true },
+              () =>
+                chatContent.shadowRoot.querySelector("assistant-message-footer")
+            );
+            let appliedButton;
+            await ContentTaskUtils.waitForCondition(() => {
+              const button = chatContent.shadowRoot
+                .querySelector("assistant-message-footer")
+                ?.shadowRoot?.querySelector("applied-memories-button");
+              // Unwrap the content element so Lit's prototype-backed
+              // properties are visible through SpecialPowers.
+              appliedButton = button?.wrappedJSObject ?? button;
+              return appliedButton?.appliedMemories?.length;
+            }, "Wait for the resume-activity memory to be marked applied");
+
+            Assert.equal(
+              appliedButton.appliedMemories.length,
+              1,
+              "Resume-activity response should have exactly one applied memory"
+            );
+            Assert.equal(
+              appliedButton.appliedMemories[0].memory_summary,
+              "Research project",
+              "Applied memory should be the one that drove the resume-activity conversation"
+            );
+          });
+        },
+        { fxAccountToken: "mock-fxa-token" }
+      );
+    } finally {
+      await SpecialPowers.popPrefEnv();
+    }
+  } finally {
+    sb.restore();
+    await stopMockOpenAI(server);
+  }
+});
+
 add_task(
   async function test_resume_prompt_click_shows_confirmation_card_without_memory_context_when_toggled_off() {
     const sb = sinon.createSandbox();

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

@@ -604,6 +604,8 @@ async function stubResumeActivityGeneration(sb, { fxAccountToken } = {}) {
     {
       id: "memory-1",
       memory_summary: "Research project",
+      lifetime_accessed_count: 0,
+      recent_accessed_counts: {},
       source_ids: {
         history_source_ids: urls
           .slice(0, 4)
@@ -613,6 +615,8 @@ async function stubResumeActivityGeneration(sb, { fxAccountToken } = {}) {
     {
       id: "memory-2",
       memory_summary: "Trip planning",
+      lifetime_accessed_count: 0,
+      recent_accessed_counts: {},
       source_ids: {
         history_source_ids: [PlacesUtils.history.hashURL(urls[4].url)],
       },
@@ -665,15 +669,20 @@ async function stubResumeActivityGenerationPool(sb, memoryCount) {
  * @param {object} sb - Sinon sandbox, owned and restored by the caller
  * @param {Function} run - Async callback invoked with
  *   {win, browser, aiWindow, buttons}
+ * @param {object} [options]
+ * @param {?string} [options.fxAccountToken] - Token for mock-server-backed
+ *   chat completion.
  */
-async function testResumeActivityClick(sb, run) {
+async function testResumeActivityClick(sb, run, { fxAccountToken } = {}) {
   await SpecialPowers.pushPrefEnv({
     set: [
       ["browser.smartwindow.memories.generateFromConversation", true],
       ["browser.smartwindow.memories.generateFromHistory", true],
     ],
   });
-  const resumeActivityStubs = await stubResumeActivityGeneration(sb);
+  const resumeActivityStubs = await stubResumeActivityGeneration(sb, {
+    fxAccountToken,
+  });
   let win;
   try {
     win = await openAIWindow();

browser/components/aiwindow/ui/test/xpcshell/test_ChatConversation.js

@@ -1010,6 +1010,131 @@ add_task(
   }
 );
 
+add_task(
+  async function test_promptEmbeddedMemories_ChatConversation_receiveResponse() {
+    let sandbox = lazy.sinon.createSandbox();
+
+    const mockMemories = [
+      { id: "mem-1", lifetime_accessed_count: 0, recent_accessed_counts: {} },
+    ];
+    sandbox.stub(MemoryStore, "getMemories").resolves(mockMemories);
+    sandbox.stub(MemoryStore, "requestSave").resolves();
+
+    const conversation = new ChatConversation({});
+    conversation.promptEmbeddedMemories = [
+      { id: "mem-1", memory_summary: "Trip planning" },
+    ];
+    conversation.addAssistantMessage("text", "some response");
+    const assistantMsg = conversation.messages.at(-1);
+
+    async function* emptyStream() {}
+    await conversation.receiveResponse(emptyStream());
+
+    Assert.deepEqual(
+      assistantMsg.memoriesApplied,
+      mockMemories,
+      "memoriesApplied should include the prompt-embedded memory, resolved through resolveUsedMemories, when nothing is cited"
+    );
+    Assert.equal(
+      mockMemories[0].lifetime_accessed_count,
+      1,
+      "Prompt-embedded memory use should be recorded like a cited one"
+    );
+    Assert.deepEqual(
+      conversation.promptEmbeddedMemories,
+      [],
+      "promptEmbeddedMemories should be consumed after the turn completes"
+    );
+
+    sandbox.restore();
+  }
+);
+
+add_task(
+  async function test_citedAndPromptEmbeddedMemoriesAreUnioned_ChatConversation_receiveResponse() {
+    let sandbox = lazy.sinon.createSandbox();
+
+    const mockMemories = [
+      { id: "mem-1", lifetime_accessed_count: 0, recent_accessed_counts: {} },
+      {
+        id: "mem-embedded",
+        lifetime_accessed_count: 0,
+        recent_accessed_counts: {},
+      },
+    ];
+    sandbox.stub(MemoryStore, "getMemories").resolves(mockMemories);
+    sandbox.stub(MemoryStore, "requestSave").resolves();
+
+    const conversation = new ChatConversation({});
+    conversation.promptEmbeddedMemories = [
+      { id: "mem-embedded", memory_summary: "Trip planning" },
+    ];
+    conversation.addAssistantMessage("text", "some response");
+    const assistantMsg = conversation.messages.at(-1);
+    assistantMsg.tokens.existing_memory = ["mem-1"];
+
+    async function* emptyStream() {}
+    await conversation.receiveResponse(emptyStream());
+
+    const memoryIds = new Set(
+      assistantMsg.memoriesApplied.map(memory => memory.id)
+    );
+    Assert.equal(
+      memoryIds.size,
+      2,
+      "memoriesApplied should include both the cited and prompt-embedded memories"
+    );
+    Assert.ok(memoryIds.has("mem-1"), "Should include the cited memory");
+    Assert.ok(
+      memoryIds.has("mem-embedded"),
+      "Should include the prompt-embedded memory"
+    );
+    Assert.deepEqual(
+      conversation.promptEmbeddedMemories,
+      [],
+      "promptEmbeddedMemories should be cleared after the turn completes"
+    );
+
+    sandbox.restore();
+  }
+);
+
+add_task(
+  async function test_citedAndPromptEmbeddedSameMemoryIsDeduped_ChatConversation_receiveResponse() {
+    let sandbox = lazy.sinon.createSandbox();
+
+    const mockMemories = [
+      { id: "mem-1", lifetime_accessed_count: 0, recent_accessed_counts: {} },
+    ];
+    sandbox.stub(MemoryStore, "getMemories").resolves(mockMemories);
+    sandbox.stub(MemoryStore, "requestSave").resolves();
+
+    const conversation = new ChatConversation({});
+    conversation.promptEmbeddedMemories = [
+      { id: "mem-1", memory_summary: "Trip planning" },
+    ];
+    conversation.addAssistantMessage("text", "some response");
+    const assistantMsg = conversation.messages.at(-1);
+    assistantMsg.tokens.existing_memory = ["mem-1"];
+
+    async function* emptyStream() {}
+    await conversation.receiveResponse(emptyStream());
+
+    Assert.equal(
+      assistantMsg.memoriesApplied.length,
+      1,
+      "The same memory cited and prompt-embedded should only appear once"
+    );
+    Assert.equal(
+      mockMemories[0].lifetime_accessed_count,
+      1,
+      "Use should be recorded once, not twice, for the same memory"
+    );
+
+    sandbox.restore();
+  }
+);
+
 add_task(function test_ChatConversation_rehydratesHistoryResultsPool() {
   const records = [
     { url: "https://example.com/1", title: "Page 1" },