tor-browser

test_ConversationSuggestions.js (43772B)

---

/* This Source Code Form is subject to the terms of the Mozilla Public
* License, v. 2.0. If a copy of the MPL was not distributed with this
* file, You can obtain one at https://mozilla.org/MPL/2.0/. */

do_get_profile();

const {
 NewTabStarterGenerator,
 trimConversation,
 addMemoriesToPrompt,
 cleanInferenceOutput,
 generateConversationStartersSidebar,
 generateFollowupPrompts,
 MemoriesGetterForSuggestionPrompts,
} = ChromeUtils.importESModule(
 "moz-src:///browser/components/aiwindow/models/ConversationSuggestions.sys.mjs"
);

const { openAIEngine } = ChromeUtils.importESModule(
 "moz-src:///browser/components/aiwindow/models/Utils.sys.mjs"
);
const { MemoriesManager } = ChromeUtils.importESModule(
 "moz-src:///browser/components/aiwindow/models/memories/MemoriesManager.sys.mjs"
);
const { MESSAGE_ROLE } = ChromeUtils.importESModule(
 "moz-src:///browser/components/aiwindow/ui/modules/ChatConstants.sys.mjs"
);
const { sinon } = ChromeUtils.importESModule(
 "resource://testing-common/Sinon.sys.mjs"
);

/**
* Constants for preference keys and test values
*/
const PREF_API_KEY = "browser.aiwindow.apiKey";
const PREF_ENDPOINT = "browser.aiwindow.endpoint";
const PREF_MODEL = "browser.aiwindow.model";
const PREF_HISTORY_ENABLED = "places.history.enabled";
const PREF_PRIVATE_BROWSING = "browser.privatebrowsing.autostart";

const API_KEY = "test-api-key";
const ENDPOINT = "https://api.test-endpoint.com/v1";
const MODEL = "test-model";

async function loadRemoteSettingsSnapshot() {
 const file = do_get_file("ai-window-prompts-remote-settings-snapshot.json");
 const data = await IOUtils.readUTF8(file.path);
 return JSON.parse(data);
}

let REAL_REMOTE_SETTINGS_SNAPSHOT;

add_setup(async function () {
 REAL_REMOTE_SETTINGS_SNAPSHOT = await loadRemoteSettingsSnapshot();
});

/**
* Cleans up preferences after testing
*/
registerCleanupFunction(() => {
 for (let pref of [
   PREF_API_KEY,
   PREF_ENDPOINT,
   PREF_MODEL,
   PREF_HISTORY_ENABLED,
   PREF_PRIVATE_BROWSING,
 ]) {
   if (Services.prefs.prefHasUserValue(pref)) {
     Services.prefs.clearUserPref(pref);
   }
 }
});

/**
* Tests for trimConversation function
*/
add_task(async function test_trimConversation() {
 const cases = [
   {
     input: [
       // empty case
     ],
     expected: [],
   },
   {
     input: [
       // more than 15 messages
       ...Array.from({ length: 20 }, (_, i) => ({
         role: i % 2 === 0 ? MESSAGE_ROLE.USER : MESSAGE_ROLE.ASSISTANT,
         content: `Message ${i + 1}`,
       })),
     ],
     expected: [
       ...Array.from({ length: 15 }, (_, i) => ({
         role: (i + 5) % 2 === 0 ? "user" : "assistant",
         content: `Message ${i + 6}`,
       })),
     ],
   },
   {
     input: [
       // should remove tool call/responses
       { role: MESSAGE_ROLE.USER, content: "What's the weather like?" },
       {
         role: MESSAGE_ROLE.ASSISTANT,
         content: "",
         tool_call: { name: "get_weather", arguments: "{}" },
       },
       {
         role: MESSAGE_ROLE.TOOL,
         content: "Here are the latest news articles.",
       },
     ],
     expected: [{ role: "user", content: "What's the weather like?" }],
   },
   {
     input: [
       // should remove system message
       { role: MESSAGE_ROLE.SYSTEM, content: "System message" },
       { role: MESSAGE_ROLE.USER, content: "Hello" },
       { role: MESSAGE_ROLE.ASSISTANT, content: "Hi there!" },
     ],
     expected: [
       { role: "user", content: "Hello" },
       { role: "assistant", content: "Hi there!" },
     ],
   },
   {
     input: [
       // should remove messages with empty content
       { role: MESSAGE_ROLE.USER, content: "\n" },
       { role: MESSAGE_ROLE.ASSISTANT, content: "   " },
     ],
     expected: [],
   },
   {
     input: [
       // no valid messages
       { role: MESSAGE_ROLE.SYSTEM, content: "System message" },
       {
         role: MESSAGE_ROLE.ASSISTANT,
         content: "",
         tool_call: { name: "get_info", arguments: "{}" },
       },
     ],
     expected: [],
   },
   {
     input: [
       // should slice after filtering invalid messages
       ...Array(10)
         .fill(0)
         .flatMap((_, i) => [
           {
             role: MESSAGE_ROLE.USER,
             content: `User message ${i + 1}`,
           },
           { role: MESSAGE_ROLE.SYSTEM, content: "System message" },
         ]),
     ],
     expected: [
       ...Array.from({ length: 10 }, (_, i) => ({
         role: "user",
         content: `User message ${i + 1}`,
       })),
     ],
   },
 ];

 for (const { input, expected } of cases) {
   const result = trimConversation(input);
   Assert.deepEqual(
     result,
     expected,
     "trimConversation should return only user/assistant messages with content"
   );
 }
});

/**
* Test for addMemoriesToPrompt function when there are memories
*/
add_task(async function test_addMemoriesToPrompt_have_memories() {
 const sb = sinon.createSandbox();
 try {
   const basePrompt = "Base prompt content.";
   const fakeMemories = ["Memory summary 1", "Memory summary 2"];
   const memoriesStub = sb
     .stub(MemoriesGetterForSuggestionPrompts, "getMemorySummariesForPrompt")
     .resolves(fakeMemories);
   const conversationMemoriesPrompt = "Memories block:\n{memories}";
   const promptWithMemories = await addMemoriesToPrompt(
     basePrompt,
     conversationMemoriesPrompt
   );

   Assert.ok(
     memoriesStub.calledOnce,
     "getMemorySummariesForPrompt should be called"
   );
   Assert.ok(
     promptWithMemories.includes("- Memory summary 1") &&
       promptWithMemories.includes("- Memory summary 2"),
     "Prompt should include memories"
   );
 } finally {
   sb.restore();
 }
});

/**
* Test for addMemoriesToPrompt function when there are no memories
*/
add_task(async function test_addMemoriesToPrompt_dont_have_memories() {
 const sb = sinon.createSandbox();
 try {
   const basePrompt = "Base prompt content.";
   const fakeMemories = [];
   const memoriesStub = sb
     .stub(MemoriesGetterForSuggestionPrompts, "getMemorySummariesForPrompt")
     .resolves(fakeMemories);
   const conversationMemoriesPrompt = "Memories block:\n{memories}";
   const promptWithMemories = await addMemoriesToPrompt(
     basePrompt,
     conversationMemoriesPrompt
   );

   Assert.ok(
     memoriesStub.calledOnce,
     "getMemorySummariesForPrompt should be called"
   );
   Assert.equal(
     promptWithMemories,
     basePrompt,
     "Prompt should be unchanged when no memories"
   );
 } finally {
   sb.restore();
 }
});

/**
* Tests for cleanInferenceOutput function
*/
add_task(function test_cleanInferenceOutput() {
 const sb = sinon.createSandbox();
 try {
   const cases = [
     {
       input: {
         finalOutput: `- Suggestion 1\n\n* Suggestion 2\n1. Suggestion 3.\n2)Suggestion 4\n[5] Suggestion 5\n6.\nLabel: Suggestion 6\nSuggestion 7.\nSuggestion 8?`,
       },
       expected: [
         "Suggestion 1",
         "Suggestion 2",
         "Suggestion 3",
         "Suggestion 4",
         "Suggestion 5",
         "Suggestion 6",
         "Suggestion 7",
         "Suggestion 8?",
       ],
     },
     {
       input: { finalOutput: `Suggestion X\nSuggestion Y\nSuggestion Z` },
       expected: ["Suggestion X", "Suggestion Y", "Suggestion Z"],
     },
     {
       input: { finalOutput: "" },
       expected: [],
     },
   ];

   for (const { input, expected } of cases) {
     const result = cleanInferenceOutput(input);
     Assert.deepEqual(
       result,
       expected,
       "cleanInferenceOutput should return expected output"
     );
   }
 } finally {
   sb.restore();
 }
});

/**
* Tests for createNewTabPromptGenerator generating prompts based on tab count
*/
add_task(
 async function test_createNewTabPromptGenerator_with_history_enabled() {
   const sb = sinon.createSandbox();
   const writingPrompts = [
     "Write a first draft",
     "Improve writing",
     "Proofread a message",
   ];

   const planningPrompts = [
     "Simplify a topic",
     "Brainstorm ideas",
     "Help make a plan",
   ];
   try {
     const cases = [
       {
         input: { tabCount: -1 },
         expectedBrowsing: [],
       },
       {
         input: { tabCount: 0 },
         expectedBrowsing: ["Find tabs in history"],
       },
       {
         input: { tabCount: 1 },
         expectedBrowsing: ["Find tabs in history", "Summarize tabs"],
       },
       {
         input: { tabCount: 2 },
         expectedBrowsing: [
           "Find tabs in history",
           "Summarize tabs",
           "Compare tabs",
         ],
       },
       {
         input: { tabCount: 3 },
         expectedBrowsing: [
           "Find tabs in history",
           "Summarize tabs",
           "Compare tabs",
         ],
       },
     ];
     const promptGenerator = NewTabStarterGenerator;
     for (const { input, expectedBrowsing } of cases) {
       const results = await promptGenerator.getPrompts(input.tabCount);
       if (input.tabCount <= -1) {
         Assert.equal(results.length, 2, "Should return 2 suggestions");
       } else {
         Assert.equal(results.length, 3, "Should return 3 suggestions");
       }
       for (const result of results) {
         Assert.equal(
           result.type,
           "chat",
           "Each result should have type 'chat'"
         );
       }
       Assert.ok(
         writingPrompts.includes(results[0].text),
         "Results should include a valid writing prompt"
       );
       Assert.ok(
         planningPrompts.includes(results[1].text),
         "Results should include a valid planning prompt"
       );
       if (results[2]) {
         Assert.ok(
           expectedBrowsing.includes(results[2].text),
           "Results should include a valid browsing prompt"
         );
       }
     }
   } finally {
     sb.restore();
   }
 }
);

/**
* Tests for createNewTabPromptGenerator generating prompts based on tab count, with history disabled
*/
add_task(
 async function test_createNewTabPromptGenerator_with_history_disabled() {
   const sb = sinon.createSandbox();
   const writingPrompts = [
     "Write a first draft",
     "Improve writing",
     "Proofread a message",
   ];

   const planningPrompts = [
     "Simplify a topic",
     "Brainstorm ideas",
     "Help make a plan",
   ];
   try {
     const cases = [
       {
         input: { tabCount: -1 },
         expectedBrowsing: [],
       },
       {
         input: { tabCount: 0 },
         expectedBrowsing: [],
       },
       {
         input: { tabCount: 1 },
         expectedBrowsing: ["Summarize tabs"],
       },
       {
         input: { tabCount: 2 },
         expectedBrowsing: ["Summarize tabs", "Compare tabs"],
       },
       {
         input: { tabCount: 3 },
         expectedBrowsing: ["Summarize tabs", "Compare tabs"],
       },
     ];
     for (const pref of [
       [{ key: PREF_HISTORY_ENABLED, value: false }],
       [{ key: PREF_PRIVATE_BROWSING, value: true }],
       [
         { key: PREF_HISTORY_ENABLED, value: false },
         { key: PREF_PRIVATE_BROWSING, value: true },
       ],
     ]) {
       for (const p of pref) {
         Services.prefs.setBoolPref(p.key, p.value);
       }
       const promptGenerator = NewTabStarterGenerator;
       for (const { input, expectedBrowsing } of cases) {
         const results = await promptGenerator.getPrompts(input.tabCount);
         if (input.tabCount <= 0) {
           Assert.equal(results.length, 2, "Should return 2 suggestions");
         } else {
           Assert.equal(results.length, 3, "Should return 3 suggestions");
         }
         for (const result of results) {
           Assert.equal(
             result.type,
             "chat",
             "Each result should have type 'chat'"
           );
         }
         Assert.ok(
           writingPrompts.includes(results[0].text),
           "Results should include a valid writing prompt"
         );
         Assert.ok(
           planningPrompts.includes(results[1].text),
           "Results should include a valid planning prompt"
         );
         if (results[2]) {
           Assert.ok(
             expectedBrowsing.includes(results[2].text),
             "Results should include a valid browsing prompt"
           );
         }
       }
       for (const p of pref) {
         Services.prefs.clearUserPref(p.key);
       }
     }
   } finally {
     sb.restore();
   }
 }
);

/**
* Tests for generateConversationStartersSidebar successfully generating suggestions
*/
add_task(async function test_generateConversationStartersSidebar_happy_path() {
 Services.prefs.setStringPref(PREF_API_KEY, API_KEY);
 Services.prefs.setStringPref(PREF_ENDPOINT, ENDPOINT);
 Services.prefs.setStringPref(PREF_MODEL, MODEL);

 const sb = sinon.createSandbox();
 try {
   // Mock the openAIEngine and memories response
   const fakeEngine = {
     run: sb.stub().resolves({
       finalOutput: `1. Suggestion 1\n\n- Suggestion 2\nLabel: Suggestion 3.\nSuggestion 4\nSuggestion 5\nSuggestion 6`,
     }),
   };
   sb.stub(openAIEngine, "_createEngine").resolves(fakeEngine);

   const fakeMemories = ["Memory summary 1", "Memory summary 2"];
   const memoriesStub = sb
     .stub(MemoriesGetterForSuggestionPrompts, "getMemorySummariesForPrompt")
     .resolves(fakeMemories);

   const n = 3;
   const contextTabs = [
     { title: "Current Tab", url: "https://current.example.com" },
     { title: "Tab 2", url: "https://tab2.example.com" },
   ];

   const result = await generateConversationStartersSidebar(
     contextTabs,
     n,
     true
   );
   Assert.ok(fakeEngine.run.calledOnce, "Engine run should be called once");
   Assert.ok(
     memoriesStub.calledOnce,
     "getMemorySummariesForPrompt should be called once"
   );

   // Verify the prompt content
   const callArgs = fakeEngine.run.firstCall.args[0];
   Assert.equal(
     callArgs.messages.length,
     2,
     "run should be called with 2 messages"
   );
   Assert.ok(
     callArgs.messages[1].content.includes(
       '{"title":"Current Tab","url":"https://current.example.com"}'
     ),
     "Prompt should include current tab info"
   );
   Assert.ok(
     callArgs.messages[1].content.includes(
       '{"title":"Tab 2","url":"https://tab2.example.com"}'
     ),
     "Prompt should include other tab info"
   );
   Assert.ok(
     callArgs.messages[1].content.includes(
       "\n- Memory summary 1\n- Memory summary 2"
     ),
     "Prompt should include memory summaries"
   );

   Assert.deepEqual(
     result,
     [
       { text: "Suggestion 1", type: "chat" },
       { text: "Suggestion 2", type: "chat" },
       { text: "Suggestion 3", type: "chat" },
     ],
     "Suggestions should match expected values"
   );
 } finally {
   sb.restore();
 }
});

/**
* Tests for generateConversationStartersSidebar without including memories
*/
add_task(
 async function test_generateConversationStartersSidebar_without_memories() {
   Services.prefs.setStringPref(PREF_API_KEY, API_KEY);
   Services.prefs.setStringPref(PREF_ENDPOINT, ENDPOINT);
   Services.prefs.setStringPref(PREF_MODEL, MODEL);

   const sb = sinon.createSandbox();
   try {
     // Mock the openAIEngine and memories response
     const fakeEngine = {
       run: sb.stub().resolves({
         finalOutput: `1. Suggestion 1\n\n- Suggestion 2\nLabel: Suggestion 3.\nSuggestion 4\nSuggestion 5\nSuggestion 6`,
       }),
     };
     sb.stub(openAIEngine, "_createEngine").resolves(fakeEngine);

     const fakeMemories = ["Memory summary 1", "Memory summary 2"];
     const memoriesStub = sb
       .stub(MemoriesGetterForSuggestionPrompts, "getMemorySummariesForPrompt")
       .resolves(fakeMemories);

     const n = 3;
     const contextTabs = [
       { title: "Current Tab", url: "https://current.example.com" },
       { title: "Tab 2", url: "https://tab2.example.com" },
     ];

     const result = await generateConversationStartersSidebar(
       contextTabs,
       n,
       false
     );
     Assert.ok(fakeEngine.run.calledOnce, "Engine run should be called once");
     Assert.ok(
       !memoriesStub.calledOnce,
       "getMemorySummariesForPrompt shouldn't be called"
     );

     // Verify the prompt content
     const callArgs = fakeEngine.run.firstCall.args[0];
     Assert.equal(
       callArgs.messages.length,
       2,
       "run should be called with 2 messages"
     );
     Assert.ok(
       callArgs.messages[1].content.includes(
         '{"title":"Current Tab","url":"https://current.example.com"}'
       ),
       "Prompt should include current tab info"
     );
     Assert.ok(
       callArgs.messages[1].content.includes(
         '{"title":"Tab 2","url":"https://tab2.example.com"}'
       ),
       "Prompt should include other tab info"
     );
     Assert.ok(
       !callArgs.messages[1].content.includes(
         "\n- Memory summary 1\n- Memory summary 2"
       ),
       "Prompt should not include memory summaries"
     );

     Assert.deepEqual(
       result,
       [
         { text: "Suggestion 1", type: "chat" },
         { text: "Suggestion 2", type: "chat" },
         { text: "Suggestion 3", type: "chat" },
       ],
       "Suggestions should match expected values"
     );
   } finally {
     sb.restore();
   }
 }
);

/**
* Tests for generateConversationStartersSidebar when no memories are returned
*/
add_task(
 async function test_generateConversationStartersSidebar_no_memories_returned() {
   Services.prefs.setStringPref(PREF_API_KEY, API_KEY);
   Services.prefs.setStringPref(PREF_ENDPOINT, ENDPOINT);
   Services.prefs.setStringPref(PREF_MODEL, MODEL);

   const sb = sinon.createSandbox();
   try {
     // Mock the openAIEngine and memories response
     const fakeEngine = {
       run: sb.stub().resolves({
         finalOutput: `1. Suggestion 1\n\n- Suggestion 2\nLabel: Suggestion 3.\nSuggestion 4\nSuggestion 5\nSuggestion 6`,
       }),
     };
     sb.stub(openAIEngine, "_createEngine").resolves(fakeEngine);

     const fakeMemories = [];
     const memoriesStub = sb
       .stub(MemoriesGetterForSuggestionPrompts, "getMemorySummariesForPrompt")
       .resolves(fakeMemories);

     const n = 3;
     const contextTabs = [
       { title: "Current Tab", url: "https://current.example.com" },
       { title: "Tab 2", url: "https://tab2.example.com" },
     ];

     const result = await generateConversationStartersSidebar(
       contextTabs,
       n,
       true
     );
     Assert.ok(fakeEngine.run.calledOnce, "Engine run should be called once");
     Assert.ok(
       memoriesStub.calledOnce,
       "getMemorySummariesForPrompt should be called once"
     );

     // Verify the prompt content
     const callArgs = fakeEngine.run.firstCall.args[0];
     Assert.equal(
       callArgs.messages.length,
       2,
       "run should be called with 2 messages"
     );
     Assert.ok(
       callArgs.messages[1].content.includes(
         '{"title":"Current Tab","url":"https://current.example.com"}'
       ),
       "Prompt should include current tab info"
     );
     Assert.ok(
       callArgs.messages[1].content.includes(
         '{"title":"Tab 2","url":"https://tab2.example.com"}'
       ),
       "Prompt should include other tab info"
     );
     Assert.ok(
       !callArgs.messages[1].content.includes("\nUser Memories:\n"),
       "Prompt shouldn't include user memories block"
     );

     Assert.deepEqual(
       result,
       [
         { text: "Suggestion 1", type: "chat" },
         { text: "Suggestion 2", type: "chat" },
         { text: "Suggestion 3", type: "chat" },
       ],
       "Suggestions should match expected values"
     );
   } finally {
     sb.restore();
   }
 }
);

/**
* Tests for generateConversationStartersSidebar when no tabs are provided
*/
add_task(async function test_generateConversationStartersSidebar_no_tabs() {
 Services.prefs.setStringPref(PREF_API_KEY, API_KEY);
 Services.prefs.setStringPref(PREF_ENDPOINT, ENDPOINT);
 Services.prefs.setStringPref(PREF_MODEL, MODEL);

 const sb = sinon.createSandbox();
 try {
   // Mock the openAIEngine and memories response
   const fakeEngine = {
     run: sb.stub().resolves({
       finalOutput: `1. Suggestion 1\n\n- Suggestion 2\nLabel: Suggestion 3.\nSuggestion 4\nSuggestion 5\nSuggestion 6`,
     }),
   };
   sb.stub(openAIEngine, "_createEngine").resolves(fakeEngine);

   const fakeMemories = ["Memory summary 1", "Memory summary 2"];
   const memoriesStub = sb
     .stub(MemoriesGetterForSuggestionPrompts, "getMemorySummariesForPrompt")
     .resolves(fakeMemories);

   const n = 3;
   const contextTabs = [];

   const result = await generateConversationStartersSidebar(
     contextTabs,
     n,
     true
   );
   Assert.ok(fakeEngine.run.calledOnce, "Engine run should be called once");
   Assert.ok(
     memoriesStub.calledOnce,
     "getMemorySummariesForPrompt should be called once"
   );

   // Verify the prompt content
   const callArgs = fakeEngine.run.firstCall.args[0];
   Assert.equal(
     callArgs.messages.length,
     2,
     "run should be called with 2 messages"
   );
   Assert.ok(
     callArgs.messages[1].content.includes("\nNo current tab\n"),
     "Prompt should indicate no current tab"
   );
   Assert.ok(
     callArgs.messages[1].content.includes("\nNo tabs available\n"),
     "Prompt should indicate no tabs available"
   );
   Assert.ok(
     callArgs.messages[1].content.includes(
       "\n- Memory summary 1\n- Memory summary 2"
     ),
     "Prompt should include memory summaries"
   );

   Assert.deepEqual(
     result,
     [
       { text: "Suggestion 1", type: "chat" },
       { text: "Suggestion 2", type: "chat" },
       { text: "Suggestion 3", type: "chat" },
     ],
     "Suggestions should match expected values"
   );
 } finally {
   sb.restore();
 }
});

/**
* Tests for generateConversationStartersSidebar with one tab provided
*/
add_task(async function test_generateConversationStartersSidebar_one_tab() {
 Services.prefs.setStringPref(PREF_API_KEY, API_KEY);
 Services.prefs.setStringPref(PREF_ENDPOINT, ENDPOINT);
 Services.prefs.setStringPref(PREF_MODEL, MODEL);

 const sb = sinon.createSandbox();
 try {
   // Mock the openAIEngine and memories response
   const fakeEngine = {
     run: sb.stub().resolves({
       finalOutput: `1. Suggestion 1\n\n- Suggestion 2\nLabel: Suggestion 3.\nSuggestion 4\nSuggestion 5\nSuggestion 6`,
     }),
   };
   sb.stub(openAIEngine, "_createEngine").resolves(fakeEngine);

   const fakeMemories = ["Memory summary 1", "Memory summary 2"];
   const memoriesStub = sb
     .stub(MemoriesGetterForSuggestionPrompts, "getMemorySummariesForPrompt")
     .resolves(fakeMemories);

   const n = 3;
   const contextTabs = [
     { title: "Only Tab", url: "https://only.example.com" },
   ];

   const result = await generateConversationStartersSidebar(
     contextTabs,
     n,
     true
   );
   Assert.ok(fakeEngine.run.calledOnce, "Engine run should be called once");
   Assert.ok(
     memoriesStub.calledOnce,
     "getMemorySummariesForPrompt should be called once"
   );

   // Verify the prompt content
   const callArgs = fakeEngine.run.firstCall.args[0];
   Assert.equal(
     callArgs.messages.length,
     2,
     "run should be called with 2 messages"
   );
   Assert.ok(
     callArgs.messages[1].content.includes(
       '\n{"title":"Only Tab","url":"https://only.example.com"}'
     ),
     "Prompt should include current tab info"
   );
   Assert.ok(
     callArgs.messages[1].content.includes("\nOnly current tab is open\n"),
     "Prompt should indicate only current tab is open"
   );
   Assert.ok(
     callArgs.messages[1].content.includes(
       "\n- Memory summary 1\n- Memory summary 2"
     ),
     "Prompt should include memory summaries"
   );

   Assert.deepEqual(
     result,
     [
       { text: "Suggestion 1", type: "chat" },
       { text: "Suggestion 2", type: "chat" },
       { text: "Suggestion 3", type: "chat" },
     ],
     "Suggestions should match expected values"
   );
 } finally {
   sb.restore();
 }
});

/**
* Tests that generateConversationStartersSidebar handles engine errors gracefully
*/
add_task(
 async function test_generateConversationStartersSidebar_engine_error() {
   Services.prefs.setStringPref(PREF_API_KEY, API_KEY);
   Services.prefs.setStringPref(PREF_ENDPOINT, ENDPOINT);
   Services.prefs.setStringPref(PREF_MODEL, MODEL);

   const sb = sinon.createSandbox();
   try {
     // Mock the openAIEngine and memories response
     const fakeEngine = {
       run: sb.stub().rejects(new Error("Engine failure")),
     };
     sb.stub(openAIEngine, "_createEngine").resolves(fakeEngine);

     const fakeMemories = ["Memory summary 1", "Memory summary 2"];
     sb.stub(
       MemoriesGetterForSuggestionPrompts,
       "getMemorySummariesForPrompt"
     ).resolves(fakeMemories);

     const n = 3;
     const contextTabs = [
       { title: "Only Tab", url: "https://only.example.com" },
     ];

     const result = await generateConversationStartersSidebar(
       contextTabs,
       n,
       true
     );
     Assert.deepEqual(result, [], "Should return empty array on engine error");
   } finally {
     sb.restore();
   }
 }
);

/**
* Tests that assistant limitations are included in conversation starter prompts
*/
add_task(
 async function test_generateConversationStartersSidebar_includes_assistant_limitations() {
   Services.prefs.setStringPref(PREF_API_KEY, API_KEY);
   Services.prefs.setStringPref(PREF_ENDPOINT, ENDPOINT);
   Services.prefs.setStringPref(PREF_MODEL, MODEL);

   const sb = sinon.createSandbox();
   try {
     const fakeEngine = {
       run: sb.stub().resolves({
         finalOutput: `Suggestion 1\nSuggestion 2\nSuggestion 3`,
       }),
     };
     sb.stub(openAIEngine, "_createEngine").resolves(fakeEngine);

     sb.stub(openAIEngine, "getRemoteClient").returns({
       get: sb.stub().resolves(REAL_REMOTE_SETTINGS_SNAPSHOT),
     });

     sb.stub(
       MemoriesGetterForSuggestionPrompts,
       "getMemorySummariesForPrompt"
     ).resolves([]);

     const n = 3;
     const contextTabs = [
       { title: "Test Tab", url: "https://test.example.com" },
     ];

     await generateConversationStartersSidebar(contextTabs, n, false);

     Assert.ok(fakeEngine.run.calledOnce, "Engine run should be called once");

     const callArgs = fakeEngine.run.firstCall.args[0];
     Assert.ok(
       callArgs.messages[1].content.includes(
         "You can do this and cannot do that."
       ),
       "Prompt should include assistant limitations from remote settings"
     );
   } finally {
     sb.restore();
   }
 }
);

/**
* Tests for generateFollowupPrompts successfully generating suggestions
*/
add_task(async function test_generateFollowupPrompts_happy_path() {
 Services.prefs.setStringPref(PREF_API_KEY, API_KEY);
 Services.prefs.setStringPref(PREF_ENDPOINT, ENDPOINT);
 Services.prefs.setStringPref(PREF_MODEL, MODEL);

 const sb = sinon.createSandbox();
 try {
   // Mock the openAIEngine and memories response
   const fakeEngine = {
     run: sb.stub().resolves({
       finalOutput: `1. Suggestion 1\n\n- Suggestion 2.\nSuggestion 3.\nSuggestion 4\nSuggestion 5\nSuggestion 6`,
     }),
   };
   sb.stub(openAIEngine, "_createEngine").resolves(fakeEngine);

   const fakeMemories = ["Memory summary 1", "Memory summary 2"];
   const memoriesStub = sb
     .stub(MemoriesGetterForSuggestionPrompts, "getMemorySummariesForPrompt")
     .resolves(fakeMemories);

   const n = 2;
   const conversationHistory = [
     { role: MESSAGE_ROLE.USER, content: "Hello" },
     { role: MESSAGE_ROLE.ASSISTANT, content: "Hi there!" },
   ];
   const currentTab = {
     title: "Current Tab",
     url: "https://current.example.com",
   };

   // Using memories
   const result = await generateFollowupPrompts(
     conversationHistory,
     currentTab,
     n,
     true
   );
   Assert.ok(fakeEngine.run.calledOnce, "Engine run should be called once");
   Assert.ok(
     memoriesStub.calledOnce,
     "getMemorySummariesForPrompt should be called once"
   );

   const callArgs = fakeEngine.run.firstCall.args[0];
   Assert.equal(
     callArgs.messages.length,
     2,
     "run should be called with 2 messages"
   );
   Assert.ok(
     callArgs.messages[1].content.includes(
       '{"title":"Current Tab","url":"https://current.example.com"}'
     ),
     "Prompt should include current tab info"
   );
   Assert.ok(
     callArgs.messages[1].content.includes(
       '[{"role":"user","content":"Hello"},{"role":"assistant","content":"Hi there!"}]'
     ),
     "Prompt should include conversation history"
   );
   Assert.ok(
     callArgs.messages[1].content.includes(
       "\n- Memory summary 1\n- Memory summary 2"
     ),
     "Prompt should include memory summaries"
   );

   Assert.deepEqual(
     result,
     [
       { text: "Suggestion 1", type: "chat" },
       { text: "Suggestion 2", type: "chat" },
     ],
     "Suggestions should match"
   );
 } finally {
   sb.restore();
 }
});

/**
* Tests for generateFollowupPrompts without including memories
*/
add_task(async function test_generateFollowupPrompts_no_memories() {
 Services.prefs.setStringPref(PREF_API_KEY, API_KEY);
 Services.prefs.setStringPref(PREF_ENDPOINT, ENDPOINT);
 Services.prefs.setStringPref(PREF_MODEL, MODEL);

 const sb = sinon.createSandbox();
 try {
   // Mock the openAIEngine and memories response
   const fakeEngine = {
     run: sb.stub().resolves({
       finalOutput: `1. Suggestion 1\n\n- Suggestion 2.\nSuggestion 3.\nSuggestion 4\nSuggestion 5\nSuggestion 6`,
     }),
   };
   sb.stub(openAIEngine, "_createEngine").resolves(fakeEngine);

   const fakeMemories = ["Memory summary 1", "Memory summary 2"];
   const memoriesStub = sb
     .stub(MemoriesGetterForSuggestionPrompts, "getMemorySummariesForPrompt")
     .resolves(fakeMemories);

   const n = 2;
   const conversationHistory = [
     { role: MESSAGE_ROLE.USER, content: "Hello" },
     { role: MESSAGE_ROLE.ASSISTANT, content: "Hi there!" },
   ];
   const currentTab = {
     title: "Current Tab",
     url: "https://current.example.com",
   };

   const result = await generateFollowupPrompts(
     conversationHistory,
     currentTab,
     n,
     false
   );
   Assert.ok(fakeEngine.run.calledOnce, "Engine run should be called once");
   Assert.ok(
     !memoriesStub.calledOnce,
     "getMemorySummariesForPrompt shouldn't be called"
   );

   const callArgs = fakeEngine.run.firstCall.args[0];
   Assert.equal(
     callArgs.messages.length,
     2,
     "run should be called with 2 messages"
   );
   Assert.ok(
     callArgs.messages[1].content.includes(
       '{"title":"Current Tab","url":"https://current.example.com"}'
     ),
     "Prompt should include current tab info"
   );
   Assert.ok(
     callArgs.messages[1].content.includes(
       '[{"role":"user","content":"Hello"},{"role":"assistant","content":"Hi there!"}]'
     ),
     "Prompt should include conversation history"
   );
   Assert.ok(
     !callArgs.messages[1].content.includes(
       "\n- Memory summary 1\n- Memory summary 2"
     ),
     "Prompt shouldn't include memory summaries"
   );

   Assert.deepEqual(
     result,
     [
       { text: "Suggestion 1", type: "chat" },
       { text: "Suggestion 2", type: "chat" },
     ],
     "Suggestions should match"
   );
 } finally {
   sb.restore();
 }
});

/**
* Tests for generateFollowupPrompts when no memories are returned
*/
add_task(async function test_generateFollowupPrompts_no_memories_returned() {
 Services.prefs.setStringPref(PREF_API_KEY, API_KEY);
 Services.prefs.setStringPref(PREF_ENDPOINT, ENDPOINT);
 Services.prefs.setStringPref(PREF_MODEL, MODEL);

 const sb = sinon.createSandbox();
 try {
   // Mock the openAIEngine and memories response
   const fakeEngine = {
     run: sb.stub().resolves({
       finalOutput: `1. Suggestion 1\n\n- Suggestion 2.\nSuggestion 3.\nSuggestion 4\nSuggestion 5\nSuggestion 6`,
     }),
   };
   sb.stub(openAIEngine, "_createEngine").resolves(fakeEngine);

   const fakeMemories = [];
   const memoriesStub = sb
     .stub(MemoriesGetterForSuggestionPrompts, "getMemorySummariesForPrompt")
     .resolves(fakeMemories);

   const n = 2;
   const conversationHistory = [
     { role: MESSAGE_ROLE.USER, content: "Hello" },
     { role: MESSAGE_ROLE.ASSISTANT, content: "Hi there!" },
   ];
   const currentTab = {
     title: "Current Tab",
     url: "https://current.example.com",
   };

   // Using memories
   const result = await generateFollowupPrompts(
     conversationHistory,
     currentTab,
     n,
     true
   );
   Assert.ok(fakeEngine.run.calledOnce, "Engine run should be called once");
   Assert.ok(
     memoriesStub.calledOnce,
     "getMemorySummariesForPrompt should be called once"
   );

   const callArgs = fakeEngine.run.firstCall.args[0];
   Assert.equal(
     callArgs.messages.length,
     2,
     "run should be called with 2 messages"
   );
   Assert.ok(
     callArgs.messages[1].content.includes(
       '{"title":"Current Tab","url":"https://current.example.com"}'
     ),
     "Prompt should include current tab info"
   );
   Assert.ok(
     callArgs.messages[1].content.includes(
       '[{"role":"user","content":"Hello"},{"role":"assistant","content":"Hi there!"}]'
     ),
     "Prompt should include conversation history"
   );
   Assert.ok(
     !callArgs.messages[1].content.includes("\nUser Memories:\n"),
     "Prompt shouldn't include user memories block"
   );

   Assert.deepEqual(
     result,
     [
       { text: "Suggestion 1", type: "chat" },
       { text: "Suggestion 2", type: "chat" },
     ],
     "Suggestions should match"
   );
 } finally {
   sb.restore();
 }
});

/**
* Tests for generateFollowupPrompts without a current tab
*/
add_task(async function test_generateFollowupPrompts_no_current_tab() {
 Services.prefs.setStringPref(PREF_API_KEY, API_KEY);
 Services.prefs.setStringPref(PREF_ENDPOINT, ENDPOINT);
 Services.prefs.setStringPref(PREF_MODEL, MODEL);

 const sb = sinon.createSandbox();
 try {
   // Mock the openAIEngine and memories response
   const fakeEngine = {
     run: sb.stub().resolves({
       finalOutput: `1. Suggestion 1\n\n- Suggestion 2.\nSuggestion 3.\nSuggestion 4\nSuggestion 5\nSuggestion 6`,
     }),
   };
   sb.stub(openAIEngine, "_createEngine").resolves(fakeEngine);

   const fakeMemories = [];
   const memoriesStub = sb
     .stub(MemoriesGetterForSuggestionPrompts, "getMemorySummariesForPrompt")
     .resolves(fakeMemories);

   const n = 2;
   const conversationHistory = [
     { role: MESSAGE_ROLE.USER, content: "Hello" },
     { role: MESSAGE_ROLE.ASSISTANT, content: "Hi there!" },
   ];
   const currentTab = {};

   const result = await generateFollowupPrompts(
     conversationHistory,
     currentTab,
     n,
     false
   );
   Assert.ok(fakeEngine.run.calledOnce, "Engine run should be called once");
   Assert.ok(
     !memoriesStub.calledOnce,
     "getMemorySummariesForPrompt shouldn't be called"
   );

   const callArgs = fakeEngine.run.firstCall.args[0];
   Assert.equal(
     callArgs.messages.length,
     2,
     "run should be called with 2 messages"
   );
   Assert.ok(
     callArgs.messages[1].content.includes("\nNo tab\n"),
     "Prompt shouldn't include any tab info"
   );
   Assert.ok(
     callArgs.messages[1].content.includes(
       '[{"role":"user","content":"Hello"},{"role":"assistant","content":"Hi there!"}]'
     ),
     "Prompt should include conversation history"
   );
   Assert.ok(
     !callArgs.messages[1].content.includes("\nUser Memories:\n"),
     "Prompt shouldn't include user memories block"
   );

   Assert.deepEqual(
     result,
     [
       { text: "Suggestion 1", type: "chat" },
       { text: "Suggestion 2", type: "chat" },
     ],
     "Suggestions should match"
   );
 } finally {
   sb.restore();
 }
});

/**
* Tests that generateFollowupPrompts handles engine errors gracefully
*/
add_task(async function test_generateFollowupPrompts_engine_error() {
 Services.prefs.setStringPref(PREF_API_KEY, API_KEY);
 Services.prefs.setStringPref(PREF_ENDPOINT, ENDPOINT);
 Services.prefs.setStringPref(PREF_MODEL, MODEL);

 const sb = sinon.createSandbox();
 try {
   // Mock the openAIEngine and memories response
   const fakeEngine = {
     run: sb.stub().rejects(new Error("Engine failure")),
   };
   sb.stub(openAIEngine, "_createEngine").resolves(fakeEngine);

   const fakeMemories = [];
   sb.stub(
     MemoriesGetterForSuggestionPrompts,
     "getMemorySummariesForPrompt"
   ).resolves(fakeMemories);

   const n = 2;
   const conversationHistory = [
     { role: MESSAGE_ROLE.USER, content: "Hello" },
     { role: MESSAGE_ROLE.ASSISTANT, content: "Hi there!" },
   ];
   const currentTab = {};

   const result = await generateFollowupPrompts(
     conversationHistory,
     currentTab,
     n,
     false
   );
   Assert.deepEqual(result, [], "Should return empty array on engine error");
 } finally {
   sb.restore();
 }
});

/**
* Tests that assistant limitations are included in followup prompts
*/
add_task(
 async function test_generateFollowupPrompts_includes_assistant_limitations() {
   Services.prefs.setStringPref(PREF_API_KEY, API_KEY);
   Services.prefs.setStringPref(PREF_ENDPOINT, ENDPOINT);
   Services.prefs.setStringPref(PREF_MODEL, MODEL);

   const sb = sinon.createSandbox();
   try {
     const fakeEngine = {
       run: sb.stub().resolves({
         finalOutput: `Suggestion 1\nSuggestion 2`,
       }),
     };
     sb.stub(openAIEngine, "_createEngine").resolves(fakeEngine);

     sb.stub(openAIEngine, "getRemoteClient").returns({
       get: sb.stub().resolves(REAL_REMOTE_SETTINGS_SNAPSHOT),
     });

     sb.stub(
       MemoriesGetterForSuggestionPrompts,
       "getMemorySummariesForPrompt"
     ).resolves([]);

     const n = 2;
     const conversationHistory = [
       { role: MESSAGE_ROLE.USER, content: "Hello" },
       { role: MESSAGE_ROLE.ASSISTANT, content: "Hi there!" },
     ];
     const currentTab = { title: "Test", url: "https://test.example.com" };

     await generateFollowupPrompts(conversationHistory, currentTab, n, false);

     Assert.ok(fakeEngine.run.calledOnce, "Engine run should be called once");

     const callArgs = fakeEngine.run.firstCall.args[0];
     Assert.ok(
       callArgs.messages[1].content.includes(
         "You can do this and cannot do that."
       ),
       "Prompt should include assistant limitations from remote settings"
     );
   } finally {
     sb.restore();
   }
 }
);

/**
* Tests for getMemorySummariesForPrompt happy path
*/
add_task(async function test_getMemorySummariesForPrompt_happy_path() {
 const sb = sinon.createSandbox();
 try {
   // Mock the MemoryStore to return fixed memories
   const fakeMemories = [
     {
       memory_summary: "Memory summary 1",
     },
     {
       memory_summary: "Memory summary 2",
     },
     {
       memory_summary: "Memory summary 3",
     },
   ];

   sb.stub(MemoriesManager, "getAllMemories").resolves(fakeMemories);

   const maxMemories = 2;
   const summaries =
     await MemoriesGetterForSuggestionPrompts.getMemorySummariesForPrompt(
       maxMemories
     );

   Assert.deepEqual(
     summaries,
     ["Memory summary 1", "Memory summary 2"],
     "Memory summaries should match expected values"
   );
 } finally {
   sb.restore();
 }
});

/**
* Tests for getMemorySummariesForPrompt when no memories are returned
*/
add_task(async function test_getMemorySummariesForPrompt_no_memories() {
 const sb = sinon.createSandbox();
 try {
   // Mock the MemoryStore to return fixed memories
   const fakeMemories = [];

   sb.stub(MemoriesManager, "getAllMemories").resolves(fakeMemories);

   const maxMemories = 2;
   const summaries =
     await MemoriesGetterForSuggestionPrompts.getMemorySummariesForPrompt(
       maxMemories
     );

   Assert.equal(
     summaries.length,
     0,
     `getMemorySummariesForPrompt(${maxMemories}) should return 0 summaries`
   );
 } finally {
   sb.restore();
 }
});

/**
* Tests for getMemorySummariesForPrompt with fewer memories than maxMemories
*/
add_task(async function test_getMemorySummariesForPrompt_too_few_memories() {
 const sb = sinon.createSandbox();
 try {
   // Mock the MemoryStore to return fixed memories
   const fakeMemories = [
     {
       memory_summary: "Memory summary 1",
     },
   ];

   sb.stub(MemoriesManager, "getAllMemories").resolves(fakeMemories);

   const maxMemories = 2;
   const summaries =
     await MemoriesGetterForSuggestionPrompts.getMemorySummariesForPrompt(
       maxMemories
     );

   Assert.deepEqual(
     summaries,
     ["Memory summary 1"],
     "Memory summaries should match expected values"
   );
 } finally {
   sb.restore();
 }
});

/**
* Tests for getMemorySummariesForPrompt handling duplicate summaries
*/
add_task(async function test_getMemorySummariesForPrompt_duplicates() {
 const sb = sinon.createSandbox();
 try {
   // Mock the MemoryStore to return fixed memories
   const fakeMemories = [
     {
       memory_summary: "Duplicate summary",
     },
     {
       memory_summary: "duplicate summary",
     },
     {
       memory_summary: "Unique summary",
     },
   ];

   sb.stub(MemoriesManager, "getAllMemories").resolves(fakeMemories);

   const maxMemories = 2;
   const summaries =
     await MemoriesGetterForSuggestionPrompts.getMemorySummariesForPrompt(
       maxMemories
     );

   Assert.deepEqual(
     summaries,
     ["Duplicate summary", "Unique summary"],
     "Memory summaries should match expected values"
   );
 } finally {
   sb.restore();
 }
});

/**
* Tests for getMemorySummariesForPrompt handling empty and whitespace-only summaries
*/
add_task(
 async function test_getMemorySummariesForPrompt_empty_and_whitespace() {
   const sb = sinon.createSandbox();
   try {
     // Mock the MemoryStore to return fixed memories
     const fakeMemories = [
       {
         memory_summary: "   \n",
       },
       {
         memory_summary: "",
       },
       {
         memory_summary: "Valid summary",
       },
     ];

     sb.stub(MemoriesManager, "getAllMemories").resolves(fakeMemories);

     const maxMemories = 2;
     const summaries =
       await MemoriesGetterForSuggestionPrompts.getMemorySummariesForPrompt(
         maxMemories
       );

     Assert.deepEqual(
       summaries,
       ["Valid summary"],
       "Memory summaries should match expected values"
     );
   } finally {
     sb.restore();
   }
 }
);