From 4204e30a2b374e90ef21507efa0b3da134aaf21a Mon Sep 17 00:00:00 2001 From: Elissa Cha Date: Thu, 10 Sep 2026 18:10:25 +0000 Subject: [PATCH] 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 --- .../models/ConversationSuggestions.sys.mjs | 1 + .../ai-action-confirmation.css | 2 + .../applied-memories-button.css | 1 + .../ai-website-confirmation.css | 2 + .../ui/modules/ChatConversation.sys.mjs | 28 +++- .../browser/browser_smartwindow_prompts.js | 66 +++++++++ .../aiwindow/ui/test/browser/head.js | 13 +- .../ui/test/xpcshell/test_ChatConversation.js | 125 ++++++++++++++++++ 8 files changed, 231 insertions(+), 7 deletions(-) diff --git a/browser/components/aiwindow/models/ConversationSuggestions.sys.mjs b/browser/components/aiwindow/models/ConversationSuggestions.sys.mjs index 97fd09d9fa33..165a67894edb 100644 --- a/browser/components/aiwindow/models/ConversationSuggestions.sys.mjs +++ b/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 }, diff --git a/browser/components/aiwindow/ui/components/ai-action-confirmation/ai-action-confirmation.css b/browser/components/aiwindow/ui/components/ai-action-confirmation/ai-action-confirmation.css index ea7d7a423666..b4aec52c392f 100644 --- a/browser/components/aiwindow/ui/components/ai-action-confirmation/ai-action-confirmation.css +++ b/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))); diff --git a/browser/components/aiwindow/ui/components/ai-chat-content/chat-assistant-footer/applied-memories-button.css b/browser/components/aiwindow/ui/components/ai-chat-content/chat-assistant-footer/applied-memories-button.css index 60c2e55c2308..3c7e5db52a91 100644 --- a/browser/components/aiwindow/ui/components/ai-chat-content/chat-assistant-footer/applied-memories-button.css +++ b/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%; diff --git a/browser/components/aiwindow/ui/components/ai-website-confirmation/ai-website-confirmation.css b/browser/components/aiwindow/ui/components/ai-website-confirmation/ai-website-confirmation.css index 4904555387f5..849934d886a7 100644 --- a/browser/components/aiwindow/ui/components/ai-website-confirmation/ai-website-confirmation.css +++ b/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)); } diff --git a/browser/components/aiwindow/ui/modules/ChatConversation.sys.mjs b/browser/components/aiwindow/ui/modules/ChatConversation.sys.mjs index 4665206bc29a..9f37a1dc8e7e 100644 --- a/browser/components/aiwindow/ui/modules/ChatConversation.sys.mjs +++ b/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} + */ + 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 diff --git a/browser/components/aiwindow/ui/test/browser/browser_smartwindow_prompts.js b/browser/components/aiwindow/ui/test/browser/browser_smartwindow_prompts.js index b182134d8f29..abd0c6914287 100644 --- a/browser/components/aiwindow/ui/test/browser/browser_smartwindow_prompts.js +++ b/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(); diff --git a/browser/components/aiwindow/ui/test/browser/head.js b/browser/components/aiwindow/ui/test/browser/head.js index 8c9a6d60182c..d20beb570272 100644 --- a/browser/components/aiwindow/ui/test/browser/head.js +++ b/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(); diff --git a/browser/components/aiwindow/ui/test/xpcshell/test_ChatConversation.js b/browser/components/aiwindow/ui/test/xpcshell/test_ChatConversation.js index 48d4b7159d8a..b30c6e04d3c1 100644 --- a/browser/components/aiwindow/ui/test/xpcshell/test_ChatConversation.js +++ b/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" },