Bug 2053939 - Support reopening browsing context from Pickup Where You Left Off r=pdahiya (9c3922cb) · Commits · The Tor Project / Applications / Tor Browser

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

+173 −7

Changes for browser/components/aiwindow/ui/components/ai-window/ai-window.mjs: 173 added lines, 7 removed lines.
Original line number	Diff line number	Diff line
		@@ -62,6 +62,8 @@ ChromeUtils.defineESModuleGetters(lazy, {
		"moz-src:///browser/components/aiwindow/models/ConversationSuggestions.sys.mjs",
		MAX_NUM_MEMORIES_FOR_RESUME_ACTIVITY:
		"moz-src:///browser/components/aiwindow/models/ConversationSuggestions.sys.mjs",
		constructConversationToResumeActivity:
		"moz-src:///browser/components/aiwindow/models/ConversationSuggestions.sys.mjs",
		MemoriesManager:
		"moz-src:///browser/components/aiwindow/models/memories/MemoriesManager.sys.mjs",
		getAllModelsData:
		@@ -259,6 +261,7 @@ export class AIWindow extends MozLitElement {
		#canLoadResumeStarters = true;
		#selectedModelChoiceId = null;
		#hasModelChoiceOverride = false;
		#isGeneratingResumeActivityConversation = false;

		get #kitMention() {
		return this.shadowRoot?.querySelector("kit-mention");
		@@ -1962,15 +1965,119 @@ export class AIWindow extends MozLitElement {
		};

		#handleResumePromptSelected(resumePrompt) {
		const contextMentions = resumePrompt.content.previewTabs.map(
		({ url, title }) => ({
		type: "tab",
		if (this.#isGeneratingResumeActivityConversation) {
		return;
		}

		Glean.smartWindow.quickPromptClicked.record({
		location: this.mode,
		chat_id: this.conversationId,
		message_seq: this.#conversation?.messageCount ?? 0,
		starter: true,
		});

		this.#generateResumeActivityConversation(resumePrompt).catch(e =>
		lazy.log.error("[Prompts] Resume-activity generation failed:", e)
		);
		}

		/**
		* Builds a conversation seeded with the selected resume-activity
		* suggestion and generates its first response, attaching an open_tabs
		* confirmation card built from the pill's preview tabs. Uses a plain
		* (memory-free) conversation instead of the memory-summarizing one when
		* memories are toggled off, or if building the memory-driven version
		* fails - the confirmation card itself only needs the preview tabs, not
		* the memory content.
		*
		* @param {object} resumePrompt - The clicked resume-activity prompt,
		* carrying the raw memory/content the generator needs.
		*/
		async #generateResumeActivityConversation(resumePrompt) {
		const conversationAtClick = this.#conversation;
		let conversation = null;
		this.#isGeneratingResumeActivityConversation = true;
		try {
		if (this.#memoriesToggled ?? this.#memoriesIconShown) {
		try {
		conversation = await lazy.constructConversationToResumeActivity({
		memory: resumePrompt.memory,
		content: resumePrompt.content,
		});
		} catch (e) {
		lazy.log.error(
		"[Prompts] Failed to create resume-activity conversation:",
		e
		);
		}
		}
		if (!conversation) {
		try {
		conversation = await this.#buildPlainResumeConversation(resumePrompt);
		} catch (e) {
		lazy.log.error(
		"[Prompts] Failed to create plain resume-activity conversation:",
		e
		);
		}
		}
		} finally {
		this.#isGeneratingResumeActivityConversation = false;
		}

		// The conversation may have changed while generation was in flight.
		if (
		!conversation \|\|
		this.#conversation !== conversationAtClick \|\|
		!resumePrompt.content.previewTabs?.length
		) {
		return;
		}

		const tabs = resumePrompt.content.previewTabs.map(
		({ url, title }, index) => ({
		token: String(index),
		url,
		label: title,
		title,
		iconSrc: `page-icon:${url}`,
		checked: false,
		})
		);
		this.onQuickPromptClicked(resumePrompt.text, true, contextMentions);
		await this.reloadAndGenerate(conversation, {
		uiType: "tab-group-confirmation",
		toolCallId: `resume-activity-${resumePrompt.memory.id}`,
		properties: {
		actionType: "open_tabs",
		tabGroupLabel: resumePrompt.text,
		tabs,
		},
		});
		}

		/**
		* Builds a conversation seeded with just the resume-activity headline and
		* no memory content in the system prompt - used when memories are
		* toggled off, or as a fallback if the memory-driven builder fails.
		*
		* @param {object} resumePrompt
		* @returns {Promise<ChatConversation>}
		*/
		async #buildPlainResumeConversation(resumePrompt) {
		const { engine, parameters } = await lazy.buildEngineForFeature(
		lazy.MODEL_FEATURES.CHAT,
		{ flowId: null, modelChoiceIdOverride: this.#selectedModelChoiceId }
		);
		const conversation = new lazy.ChatConversation({
		title: resumePrompt.content.headline,
		});
		conversation.engine = engine;
		conversation.parameters = parameters;
		await conversation.loadSystemPrompt();
		conversation.addUserMessage(resumePrompt.content.headline);
		conversation.securityProperties.setPrivateData();
		conversation.securityProperties.setUntrustedInput();
		conversation.securityProperties.commit();
		return conversation;
		}

		/**
		@@ -2156,10 +2263,31 @@ export class AIWindow extends MozLitElement {
		* message, or null if the user removed page context.
		* @param {boolean} [options.isRetry=false] - True when the call originated
		* from a user-initiated retry; surfaced in model_response telemetry.
		* @param {boolean} [options.ensureAssistantResponse=false] - When true and
		* no inputText is provided, add the empty assistant message the stream
		* writes into. Used when the user turn is already present in the
		* conversation (e.g. resume-activity starters), so generatePrompt is not
		* called to create it.
		* @param {boolean} [options.skipSystemPromptRefresh=false] - When true,
		* skip reloading the system prompt for this call. Used when a caller has
		* just built a bespoke system prompt (e.g. resume-activity starters) that
		* should survive this one request unchanged.
		* @param {object} [options.assistantToolUIData] - When set alongside
		* ensureAssistantResponse, attached to the empty assistant message so it
		* renders alongside the streamed response (e.g. a tab-selection card for
		* resume-activity starters).
		*/
		async #fetchAIResponse(
		inputText,
		{ skipUserDispatch = false, pageUrl, isRetry = false, ...userOpts } = {}
		{
		skipUserDispatch = false,
		pageUrl,
		isRetry = false,
		ensureAssistantResponse = false,
		skipSystemPromptRefresh = false,
		assistantToolUIData,
		...userOpts
		} = {}
		) {
		// Capture conversation and browsingContext at call time so that a tab switch
		// mid-stream cannot redirect this request to the wrong target.
		@@ -2210,8 +2338,12 @@ export class AIWindow extends MozLitElement {

		// Rewrites the system prompt in place so a restored conversation gets
		// today's timestamp and the latest RS content. The engine was just built
		// for this model choice, so its model drives the v2 assembly.
		// for this model choice, so its model drives the v2 assembly. Skipped
		// when a caller has just built a bespoke system prompt of its own,
		// so it survives this one request unchanged.
		if (!skipSystemPromptRefresh) {
		await conversation.loadSystemPrompt();
		}

		if (inputText) {
		await conversation.generatePrompt(
		@@ -2223,6 +2355,14 @@ export class AIWindow extends MozLitElement {

		conversation.addAssistantMessage("text", "");

		this.#sendModelRequestTelemetryEvent();
		} else if (ensureAssistantResponse) {
		// The user turn is already in the conversation, so generatePrompt is
		// skipped; add the empty assistant message receiveResponse streams into.
		const assistantMessage = conversation.addAssistantMessage("text", "");
		if (assistantToolUIData) {
		assistantMessage.toolUIData = assistantToolUIData;
		}
		this.#sendModelRequestTelemetryEvent();
		}

		@@ -2931,6 +3071,32 @@ export class AIWindow extends MozLitElement {
		this.#continueAfterToolResult();
		}

		/**
		* Opens a pre-built conversation whose last message is a user turn and
		* generates the assistant reply. Unlike reloadAndContinue (which resumes
		* an in-flight tool turn), this adds the assistant placeholder the stream
		* needs, and skips the usual system-prompt reload since the conversation
		* was just built with its own bespoke system prompt.
		*
		* @param {ChatConversation} conversation
		* @param {object} [assistantToolUIData] - Attached to the assistant
		* message so it renders alongside the streamed response (e.g. a
		* tab-selection card for resume-activity starters).
		*/
		async reloadAndGenerate(conversation, assistantToolUIData) {
		if (!conversation) {
		return;
		}
		const userMessage = conversation.messages.at(-1);
		await conversation.injectRealTimeContext(userMessage, {});
		this.openConversation(conversation);
		this.#fetchAIResponse(undefined, {
		ensureAssistantResponse: true,
		skipSystemPromptRefresh: true,
		assistantToolUIData,
		});
		}

		async #continueAfterToolResult() {
		// Show searching indicator if the last tool was run_search
		const lastToolCall = this.#conversation.messages

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

+165 −17

Changes for browser/components/aiwindow/ui/test/browser/browser_smartwindow_prompts.js: 165 added lines, 17 removed lines.
Original line number	Diff line number	Diff line
		@@ -245,6 +245,41 @@ async function stubResumeActivityGeneration(sb) {
		};
		}

		/**
		* Shared setup/teardown for tests that click a resume pill: enables the
		* memories prefs, stubs resume-activity generation so a real pill renders,
		* opens the AI Window, and hands the caller its buttons to click. `run`
		* supplies whatever additional stubs it needs (engine build, fetchWithHistory,
		* etc.) via its own sandbox before calling this.
		*
		* @param {object} sb - Sinon sandbox, owned and restored by the caller
		* @param {Function} run - Async callback invoked with
		* {win, browser, aiWindow, buttons}
		*/
		async function testResumeActivityClick(sb, run) {
		await SpecialPowers.pushPrefEnv({
		set: [
		["browser.smartwindow.memories.generateFromConversation", true],
		["browser.smartwindow.memories.generateFromHistory", true],
		],
		});
		const resumeActivityStubs = await stubResumeActivityGeneration(sb);
		let win;
		try {
		win = await openAIWindow();
		const browser = win.gBrowser.selectedBrowser;
		const aiWindow = browser.contentDocument.querySelector("ai-window");
		const buttons = await getPromptButtons(browser);
		await run({ win, browser, aiWindow, buttons });
		} finally {
		if (win) {
		await BrowserTestUtils.closeWindow(win);
		}
		await resumeActivityStubs.cleanup();
		await SpecialPowers.popPrefEnv();
		}
		}

		describe("sidebar conversation starter prompts", () => {
		let responseContent, mock, gAiWindow, backgroundTab;

		@@ -585,23 +620,6 @@ add_task(async function test_fullpage_resume_starters() {
		},
		"The resume pill should retain favicons, tab payload, and memory id"
		);

		const submitChatMessageStub = sb.stub(aiWindow, "submitChatMessage");
		buttons[0].click();
		Assert.ok(
		submitChatMessageStub.calledOnce,
		"Clicking a resume pill should submit a chat message"
		);
		Assert.deepEqual(
		submitChatMessageStub.firstCall.args[0].contextMentions,
		[1, 2, 3, 4].map(id => ({
		type: "tab",
		url: `https://example.com/${id}`,
		label: `Example ${id}`,
		iconSrc: `page-icon:https://example.com/${id}`,
		})),
		"Clicking a resume pill should submit its preview tabs as context"
		);
		} finally {
		if (win) {
		await BrowserTestUtils.closeWindow(win);
		@@ -727,6 +745,136 @@ add_task(
		}
		);

		add_task(async function test_resume_prompt_click_shows_confirmation_card() {
		const sb = sinon.createSandbox();
		try {
		sb.stub(openAIEngine, "build").resolves({});
		const fetchWithHistoryStub = sb.stub(Chat, "fetchWithHistory").resolves();

		await testResumeActivityClick(sb, async ({ aiWindow, buttons }) => {
		buttons[0].click();

		await TestUtils.waitForCondition(
		() => fetchWithHistoryStub.calledOnce,
		"Should generate a response for the resume-activity conversation"
		);

		const assistantMessage = aiWindow.conversation.messages.at(-1);
		Assert.deepEqual(
		assistantMessage.toolUIData,
		{
		uiType: "tab-group-confirmation",
		toolCallId: "resume-activity-memory-1",
		properties: {
		actionType: "open_tabs",
		tabGroupLabel: "Pick up your research",
		tabs: [1, 2, 3, 4].map((id, index) => ({
		token: String(index),
		url: `https://example.com/${id}`,
		title: `Example ${id}`,
		iconSrc: `page-icon:https://example.com/${id}`,
		checked: false,
		})),
		},
		},
		"Should attach a tab-selection confirmation card built from the pill's preview tabs"
		);
		});
		} finally {
		sb.restore();
		}
		});

		add_task(async function test_resume_prompt_click_injects_context() {
		const sb = sinon.createSandbox();
		try {
		sb.stub(openAIEngine, "build").resolves({});
		const fetchWithHistoryStub = sb.stub(Chat, "fetchWithHistory").resolves();
		const realTimeStub = sb
		.stub(this.ChatConversation.prototype, "injectRealTimeContext")
		.resolves();

		await testResumeActivityClick(sb, async ({ buttons }) => {
		buttons[0].click();

		await TestUtils.waitForCondition(
		() => fetchWithHistoryStub.calledOnce,
		"Should generate a response for the resume-activity conversation"
		);

		Assert.ok(
		realTimeStub.calledOnce,
		"Should inject real-time context onto the resume-activity user message"
		);
		});
		} finally {
		sb.restore();
		}
		});

		add_task(
		async function test_resume_prompt_click_shows_confirmation_card_without_memory_context_when_toggled_off() {
		const sb = sinon.createSandbox();
		try {
		sb.stub(openAIEngine, "build").resolves({});
		const fetchWithHistoryStub = sb.stub(Chat, "fetchWithHistory").resolves();

		await testResumeActivityClick(sb, async ({ aiWindow, buttons }) => {
		const memoriesButton = aiWindow.shadowRoot.querySelector(
		"memories-icon-button"
		);
		memoriesButton.dispatchEvent(
		new CustomEvent("aiwindow-memories-toggle:on-change", {
		bubbles: true,
		composed: true,
		detail: { pressed: false },
		})
		);
		await TestUtils.waitForCondition(
		() => memoriesButton.pressed === false,
		"Memories toggle should turn off"
		);

		buttons[0].click();

		await TestUtils.waitForCondition(
		() => fetchWithHistoryStub.calledOnce,
		"Should generate a response for the resume-activity conversation"
		);

		const assistantMessage = aiWindow.conversation.messages.at(-1);
		Assert.deepEqual(
		assistantMessage.toolUIData,
		{
		uiType: "tab-group-confirmation",
		toolCallId: "resume-activity-memory-1",
		properties: {
		actionType: "open_tabs",
		tabGroupLabel: "Pick up your research",
		tabs: [1, 2, 3, 4].map((id, index) => ({
		token: String(index),
		url: `https://example.com/${id}`,
		title: `Example ${id}`,
		iconSrc: `page-icon:https://example.com/${id}`,
		checked: false,
		})),
		},
		},
		"Should still attach a tab-selection confirmation card when memories are toggled off"
		);

		const userMessage = aiWindow.conversation.messages.at(-2);
		Assert.ok(
		!userMessage.content.userContext?.resumeActivityContext,
		"Should not bake memory context into the conversation when memories are toggled off"
		);
		});
		} finally {
		sb.restore();
		}
		}
		);

		add_task(
		async function test_fullpage_resume_starters_disabled_without_existing_memories() {
		const sb = sinon.createSandbox();