tor-browser

test_ChatUtils.js (15838B)

---

/**
* 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 http://mozilla.org/MPL/2.0/.
*/

do_get_profile();

const {
 constructRealTimeInfoInjectionMessage,
 getLocalIsoTime,
 getCurrentTabMetadata,
 constructRelevantMemoriesContextMessage,
 parseContentWithTokens,
 detectTokens,
} = ChromeUtils.importESModule(
 "moz-src:///browser/components/aiwindow/models/ChatUtils.sys.mjs"
);
const { MemoriesManager } = ChromeUtils.importESModule(
 "moz-src:///browser/components/aiwindow/models/memories/MemoriesManager.sys.mjs"
);
const { MemoryStore } = ChromeUtils.importESModule(
 "moz-src:///browser/components/aiwindow/services/MemoryStore.sys.mjs"
);
const { sinon } = ChromeUtils.importESModule(
 "resource://testing-common/Sinon.sys.mjs"
);

/**
* Constants for test memories
*/
const TEST_MEMORIES = [
 {
   memory_summary: "Loves drinking coffee",
   category: "Food & Drink",
   intent: "Plan / Organize",
   score: 3,
 },
 {
   memory_summary: "Buys dog food online",
   category: "Pets & Animals",
   intent: "Buy / Acquire",
   score: 4,
 },
];

/**
* Helper function bulk-add memories
*/
async function clearAndAddMemories() {
 const memories = await MemoryStore.getMemories();
 for (const memory of memories) {
   await MemoryStore.hardDeleteMemory(memory.id);
 }
 for (const memory of TEST_MEMORIES) {
   await MemoryStore.addMemory(memory);
 }
}

/**
* 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 API_KEY = "fake-key";
const ENDPOINT = "https://api.fake-endpoint.com/v1";
const MODEL = "fake-model";

add_setup(async function () {
 // Setup prefs used across multiple tests
 Services.prefs.setStringPref(PREF_API_KEY, API_KEY);
 Services.prefs.setStringPref(PREF_ENDPOINT, ENDPOINT);
 Services.prefs.setStringPref(PREF_MODEL, MODEL);

 // Clear prefs after testing
 registerCleanupFunction(() => {
   for (let pref of [PREF_API_KEY, PREF_ENDPOINT, PREF_MODEL]) {
     if (Services.prefs.prefHasUserValue(pref)) {
       Services.prefs.clearUserPref(pref);
     }
   }
 });
});

add_task(function test_getLocalIsoTime_returns_offset_timestamp() {
 const sb = sinon.createSandbox();
 const clock = sb.useFakeTimers({ now: Date.UTC(2025, 11, 27, 14, 0, 0) });
 try {
   const iso = getLocalIsoTime();
   Assert.ok(
     typeof iso === "string" && !!iso.length,
     "Should return a non-empty string"
   );
   Assert.ok(
     /^\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}$/.test(iso),
     "Should include date, time (up to seconds), without timezone offset"
   );
 } finally {
   clock.restore();
   sb.restore();
 }
});

add_task(async function test_getCurrentTabMetadata_fetch_fallback() {
 const sb = sinon.createSandbox();
 const tracker = { getTopWindow: sb.stub() };
 const pageData = {
   getCached: sb.stub(),
 };
 const fakeActor = {
   collectPageData: sb.stub().resolves({
     description: "Collected description",
   }),
 };
 const fakeBrowser = {
   currentURI: { spec: "https://example.com/article" },
   contentTitle: "",
   documentTitle: "Example Article",
   browsingContext: {
     currentWindowGlobal: {
       getActor: sb.stub().returns(fakeActor),
     },
   },
 };

 tracker.getTopWindow.returns({
   gBrowser: { selectedBrowser: fakeBrowser },
 });
 pageData.getCached.returns(null);

 try {
   const result = await getCurrentTabMetadata({
     BrowserWindowTracker: tracker,
     PageDataService: pageData,
   });
   Assert.deepEqual(result, {
     url: "https://example.com/article",
     title: "Example Article",
     description: "Collected description",
   });
   Assert.ok(
     fakeActor.collectPageData.calledOnce,
     "Should collect page data from actor when not cached"
   );
   Assert.ok(
     fakeBrowser.browsingContext.currentWindowGlobal.getActor.calledWith(
       "PageData"
     ),
     "Should get PageData actor"
   );
 } finally {
   sb.restore();
 }
});

add_task(
 async function test_constructRealTimeInfoInjectionMessage_with_tab_info() {
   const sb = sinon.createSandbox();
   const tracker = { getTopWindow: sb.stub() };
   const pageData = {
     getCached: sb.stub(),
   };
   const locale = Services.locale.appLocaleAsBCP47;
   const fakeActor = {
     collectPageData: sb.stub(),
   };
   const fakeBrowser = {
     currentURI: { spec: "https://mozilla.org" },
     contentTitle: "Mozilla",
     documentTitle: "Mozilla",
     browsingContext: {
       currentWindowGlobal: {
         getActor: sb.stub().returns(fakeActor),
       },
     },
   };

   tracker.getTopWindow.returns({
     gBrowser: { selectedBrowser: fakeBrowser },
   });
   pageData.getCached.returns({
     description: "Internet for people",
   });
   const clock = sb.useFakeTimers({ now: Date.UTC(2025, 11, 27, 14, 0, 0) });

   try {
     const message = await constructRealTimeInfoInjectionMessage({
       BrowserWindowTracker: tracker,
       PageDataService: pageData,
     });
     Assert.equal(message.role, "system", "Should return system role");
     Assert.ok(
       message.content.includes(`Locale: ${locale}`),
       "Should include locale"
     );
     Assert.ok(
       message.content.includes("Current active browser tab details:"),
       "Should include tab details heading"
     );
     Assert.ok(
       message.content.includes("- URL: https://mozilla.org"),
       "Should include tab URL"
     );
     Assert.ok(
       message.content.includes("- Title: Mozilla"),
       "Should include tab title"
     );
     Assert.ok(
       message.content.includes("- Description: Internet for people"),
       "Should include tab description"
     );
     Assert.ok(
       fakeActor.collectPageData.notCalled,
       "Should not collect page data when cached data exists"
     );
   } finally {
     clock.restore();
     sb.restore();
   }
 }
);

add_task(
 async function test_constructRealTimeInfoInjectionMessage_without_tab_info() {
   const sb = sinon.createSandbox();
   const tracker = { getTopWindow: sb.stub() };
   const pageData = {
     getCached: sb.stub(),
     fetchPageData: sb.stub(),
   };
   const locale = Services.locale.appLocaleAsBCP47;

   tracker.getTopWindow.returns(null);
   const clock = sb.useFakeTimers({ now: Date.UTC(2025, 11, 27, 14, 0, 0) });

   try {
     const message = await constructRealTimeInfoInjectionMessage({
       BrowserWindowTracker: tracker,
       PageDataService: pageData,
     });
     Assert.ok(
       message.content.includes("No active browser tab."),
       "Should mention missing tab info"
     );
     Assert.ok(
       !message.content.includes("- URL:"),
       "Should not include empty tab fields"
     );
     Assert.ok(
       message.content.includes(`Locale: ${locale}`),
       "Should include system locale"
     );
   } finally {
     clock.restore();
     sb.restore();
   }
 }
);

add_task(async function test_constructRelevantMemoriesContextMessage() {
 await clearAndAddMemories();

 const sb = sinon.createSandbox();
 try {
   const fakeEngine = {
     run() {
       return {
         finalOutput: `{
           "categories": ["Food & Drink"],
           "intents": ["Plan / Organize"]
         }`,
       };
     },
   };

   // Stub the `ensureOpenAIEngine` method in MemoriesManager
   const stub = sb
     .stub(MemoriesManager, "ensureOpenAIEngine")
     .returns(fakeEngine);

   const relevantMemoriesContextMessage =
     await constructRelevantMemoriesContextMessage("I love drinking coffee");
   Assert.ok(stub.calledOnce, "ensureOpenAIEngine should be called once");

   // Check relevantMemoriesContextMessage's top level structure
   Assert.strictEqual(
     typeof relevantMemoriesContextMessage,
     "object",
     "Should return an object"
   );
   Assert.equal(
     Object.keys(relevantMemoriesContextMessage).length,
     2,
     "Should have 2 keys"
   );

   // Check specific fields
   Assert.equal(
     relevantMemoriesContextMessage.role,
     "system",
     "Should have role 'system'"
   );
   Assert.ok(
     typeof relevantMemoriesContextMessage.content === "string" &&
       relevantMemoriesContextMessage.content.length,
     "Content should be a non-empty string"
   );

   const content = relevantMemoriesContextMessage.content;
   Assert.ok(
     content.includes(
       "Use them to personalized your response using the following guidelines:"
     ),
     "Relevant memories context prompt should pull from the correct base"
   );
   Assert.ok(
     content.includes("- Loves drinking coffee"),
     "Content should include relevant memory"
   );
   Assert.ok(
     !content.includes("- Buys dog food online"),
     "Content should not include non-relevant memory"
   );
 } finally {
   sb.restore();
 }
});

add_task(
 async function test_constructRelevantMemoriesContextMessage_no_relevant_memories() {
   await clearAndAddMemories();

   const sb = sinon.createSandbox();
   try {
     const fakeEngine = {
       run() {
         return {
           finalOutput: `{
           "categories": ["Health & Fitness"],
           "intents": ["Plan / Organize"]
         }`,
         };
       },
     };

     // Stub the `ensureOpenAIEngine` method in MemoriesManager
     const stub = sb
       .stub(MemoriesManager, "ensureOpenAIEngine")
       .returns(fakeEngine);

     const relevantMemoriesContextMessage =
       await constructRelevantMemoriesContextMessage("I love drinking coffee");
     Assert.ok(stub.calledOnce, "ensureOpenAIEngine should be called once");

     // No relevant memories, so returned value should be null
     Assert.equal(
       relevantMemoriesContextMessage,
       null,
       "Should return null when there are no relevant memories"
     );
   } finally {
     sb.restore();
   }
 }
);

add_task(async function test_parseContentWithTokens_no_tokens() {
 const content = "This is a regular message with no special tokens.";
 const result = await parseContentWithTokens(content);

 Assert.equal(
   result.cleanContent,
   content,
   "Clean content should match original when no tokens present"
 );
 Assert.equal(result.searchQueries.length, 0, "Should have no search queries");
 Assert.equal(result.usedMemories.length, 0, "Should have no used memories");
});

add_task(async function test_parseContentWithTokens_single_search_token() {
 const content =
   "You can find great coffee in the downtown area.§search: best coffee shops near me§";
 const result = await parseContentWithTokens(content);

 Assert.equal(
   result.cleanContent,
   "You can find great coffee in the downtown area.",
   "Should remove search token from content"
 );
 Assert.equal(result.searchQueries.length, 1, "Should have one search query");
 Assert.equal(
   result.searchQueries[0],
   "best coffee shops near me",
   "Should extract correct search query"
 );
 Assert.equal(result.usedMemories.length, 0, "Should have no used memories");
});

add_task(async function test_parseContentWithTokens_single_memory_token() {
 const content =
   "I recommend trying herbal tea blends.§existing_memory: likes tea§";
 const result = await parseContentWithTokens(content);

 Assert.equal(
   result.cleanContent,
   "I recommend trying herbal tea blends.",
   "Should remove memory token from content"
 );
 Assert.equal(result.searchQueries.length, 0, "Should have no search queries");
 Assert.equal(result.usedMemories.length, 1, "Should have one used memory");
 Assert.equal(
   result.usedMemories[0],
   "likes tea",
   "Should extract correct memory"
 );
});

add_task(async function test_parseContentWithTokens_multiple_mixed_tokens() {
 const content =
   "I recommend checking out organic coffee options.§existing_memory: prefers organic§ They have great flavor profiles.§search: organic coffee beans reviews§§search: best organic cafes nearby§";
 const result = await parseContentWithTokens(content);

 Assert.equal(
   result.cleanContent,
   "I recommend checking out organic coffee options. They have great flavor profiles.",
   "Should remove all tokens from content"
 );
 Assert.equal(
   result.searchQueries.length,
   2,
   "Should have two search queries"
 );
 Assert.deepEqual(
   result.searchQueries,
   ["organic coffee beans reviews", "best organic cafes nearby"],
   "Should extract search queries in correct order"
 );
 Assert.equal(result.usedMemories.length, 1, "Should have one used memory");
 Assert.equal(
   result.usedMemories[0],
   "prefers organic",
   "Should extract correct memory"
 );
});

add_task(async function test_parseContentWithTokens_tokens_with_whitespace() {
 const content =
   "You can find more details online.§search:   coffee brewing methods   §";
 const result = await parseContentWithTokens(content);

 Assert.equal(
   result.cleanContent,
   "You can find more details online.",
   "Should remove token with whitespace"
 );
 Assert.equal(result.searchQueries.length, 1, "Should have one search query");
 Assert.equal(
   result.searchQueries[0],
   "coffee brewing methods",
   "Should trim whitespace from extracted query"
 );
});

add_task(async function test_parseContentWithTokens_adjacent_tokens() {
 const content =
   "Here are some great Italian dining options.§existing_memory: prefers italian food§§search: local italian restaurants§";
 const result = await parseContentWithTokens(content);

 Assert.equal(
   result.cleanContent,
   "Here are some great Italian dining options.",
   "Should remove adjacent tokens"
 );
 Assert.equal(result.searchQueries.length, 1, "Should have one search query");
 Assert.equal(
   result.searchQueries[0],
   "local italian restaurants",
   "Should extract search query"
 );
 Assert.equal(result.usedMemories.length, 1, "Should have one memory");
 Assert.equal(
   result.usedMemories[0],
   "prefers italian food",
   "Should extract memory"
 );
});

add_task(function test_detectTokens_basic_pattern() {
 const content =
   "There are many great options available.§search: coffee shops near downtown§§search: best rated restaurants§";
 const searchRegex = /§search:\s*([^§]+)§/gi;
 const result = detectTokens(content, searchRegex, "query");

 Assert.equal(result.length, 2, "Should find two matches");
 Assert.equal(
   result[0].query,
   "coffee shops near downtown",
   "First match should extract correct query"
 );
 Assert.equal(
   result[0].fullMatch,
   "§search: coffee shops near downtown§",
   "First match should include full match"
 );
 Assert.equal(
   result[0].startIndex,
   39,
   "First match should have correct start index"
 );
 Assert.equal(
   result[1].query,
   "best rated restaurants",
   "Second match should extract correct query"
 );
});

add_task(function test_detectTokens_custom_key() {
 const content =
   "I recommend trying the Thai curry.§memory: prefers spicy food§";
 const memoryRegex = /§memory:\s*([^§]+)§/gi;
 const result = detectTokens(content, memoryRegex, "customKey");

 Assert.equal(result.length, 1, "Should find one match");
 Assert.equal(
   result[0].customKey,
   "prefers spicy food",
   "Should use custom key for extracted value"
 );
 Assert.ok(
   result[0].hasOwnProperty("customKey"),
   "Result should have the custom key property"
 );
 Assert.ok(
   !result[0].hasOwnProperty("query"),
   "Result should not have default 'query' property"
 );
});